factor/core/optimizer/math/math.factor

! Copyright (C) 2005, 2008 Slava Pestov.
! See http://factorcode.org/license.txt for BSD license.
USING: effects alien alien.accessors arrays generic hashtables
kernel assocs math math.libm math.private kernel.private
sequences words parser vectors strings sbufs io namespaces
assocs quotations math.intervals sequences.private combinators
splitting layouts math.parser classes classes.algebra
generic.math inference.class inference.dataflow
optimizer.pattern-match optimizer.backend optimizer.def-use
optimizer.inlining optimizer.math.partial generic.standard
system accessors ;
IN: optimizer.math

: define-math-identities ( word identities -- )
    >r all-derived-ops r> define-identities ;

\ number= {
    { { @ @ } [ 2drop t ] }
} define-math-identities

\ + {
    { { number 0 } [ drop ] }
    { { 0 number } [ nip ] }
} define-math-identities

\ - {
    { { number 0 } [ drop ] }
    { { @ @ } [ 2drop 0 ] }
} define-math-identities

\ < {
    { { @ @ } [ 2drop f ] }
} define-math-identities

\ <= {
    { { @ @ } [ 2drop t ] }
} define-math-identities

\ > {
    { { @ @ } [ 2drop f ] }
} define-math-identities

\ >= {
    { { @ @ } [ 2drop t ] }
} define-math-identities

\ * {
    { { number 1 } [ drop ] }
    { { 1 number } [ nip ] }
    { { number 0 } [ nip ] }
    { { 0 number } [ drop ] }
    { { number -1 } [ drop 0 swap - ] }
    { { -1 number } [ nip 0 swap - ] }
} define-math-identities

\ / {
    { { number 1 } [ drop ] }
    { { number -1 } [ drop 0 swap - ] }
} define-math-identities

\ mod {
    { { integer 1 } [ 2drop 0 ] }
} define-math-identities

\ rem {
    { { integer 1 } [ 2drop 0 ] }
} define-math-identities

\ bitand {
    { { number -1 } [ drop ] }
    { { -1 number } [ nip ] }
    { { @ @ } [ drop ] }
    { { number 0 } [ nip ] }
    { { 0 number } [ drop ] }
} define-math-identities

\ bitor {
    { { number 0 } [ drop ] }
    { { 0 number } [ nip ] }
    { { @ @ } [ drop ] }
    { { number -1 } [ nip ] }
    { { -1 number } [ drop ] }
} define-math-identities

\ bitxor {
    { { number 0 } [ drop ] }
    { { 0 number } [ nip ] }
    { { number -1 } [ drop bitnot ] }
    { { -1 number } [ nip bitnot ] }
    { { @ @ } [ 2drop 0 ] }
} define-math-identities

\ shift {
    { { 0 number } [ drop ] }
    { { number 0 } [ drop ] }
} define-math-identities

: math-closure ( class -- newclass )
    { null fixnum bignum integer rational float real number }
    [ class<= ] with find nip number or ;

: fits? ( interval class -- ? )
    "interval" word-prop dup
    [ interval-subset? ] [ 2drop t ] if ;

: math-output-class ( node upgrades -- newclass )
    >r
    in-d>> null [ value-class* math-closure math-class-max ] reduce
    dup r> at swap or ;

: won't-overflow? ( interval node -- ? )
    node-in-d [ value-class* fixnum class<= ] all?
    swap fixnum fits? and ;

: post-process ( class interval node -- classes intervals )
    dupd won't-overflow?
    [ >r dup { f integer } member? [ drop fixnum ] when r> ] when
    [ dup [ 1array ] when ] bi@ ;

: math-output-interval-1 ( node word -- interval )
    dup [
        >r node-in-d first value-interval* dup
        [ r> execute ] [ r> 2drop f ] if
    ] [
        2drop f
    ] if ; inline

: math-output-class/interval-1 ( node word -- classes intervals )
    [ drop { } math-output-class 1array ]
    [ math-output-interval-1 1array ] 2bi ;

{
    { bitnot interval-bitnot }
    { fixnum-bitnot interval-bitnot }
    { bignum-bitnot interval-bitnot }
} [
    [ math-output-class/interval-1 ] curry
    "output-classes" set-word-prop
] assoc-each

: intervals ( node -- i1 i2 )
    node-in-d first2 [ value-interval* ] bi@ ;

: math-output-interval-2 ( node word -- interval )
    dup [
        >r intervals 2dup and [ r> execute ] [ r> 3drop f ] if
    ] [
        2drop f
    ] if ; inline

: math-output-class/interval-2 ( node upgrades word -- classes intervals )
    pick >r
    >r over r>
    math-output-interval-2
    >r math-output-class r>
    r> post-process ; inline

{
    { + { { fixnum integer } } interval+ }
    { - { { fixnum integer } } interval- }
    { * { { fixnum integer } } interval* }
    { / { { fixnum rational } { integer rational } } interval/-safe }
    { /i { { fixnum integer } } interval/i }
    { shift { { fixnum integer } } interval-shift-safe }
} [
    first3 [
        [
            math-output-class/interval-2
        ] 2curry "output-classes" set-word-prop
    ] 2curry each-derived-op
] each

: math-output-class/interval-2-fast ( node word -- classes intervals )
    math-output-interval-2 fixnum [ 1array ] bi@ swap ; inline

[
    { + interval+ }
    { - interval- }
    { * interval* }
    { shift interval-shift-safe }
] [
    first2 [
        [
            math-output-class/interval-2-fast
        ] curry "output-classes" set-word-prop
    ] curry each-fast-derived-op
] each

: real-value? ( value -- n ? )
    dup value? [ value-literal dup real? ] [ drop f f ] if ;

: mod-range ( n -- interval )
    dup neg swap (a,b) ;

: rem-range ( n -- interval )
    0 swap [a,b) ;

: bitand-range ( n -- interval )
    dup 0 < [ drop f ] [ 0 swap [a,b] ] if ;

: math-output-interval-special ( node word -- interval )
    dup [
        >r node-in-d second real-value?
        [ r> execute ] [ r> 2drop f ] if
    ] [
        2drop f
    ] if ; inline

: math-output-class/interval-special ( node min word -- classes intervals )
    pick >r
    >r over r>
    math-output-interval-special
    >r math-output-class r>
    r> post-process ; inline

{
    { mod { } mod-range }
    { rem { { fixnum integer } } rem-range }

    { bitand { } bitand-range }
    { bitor { } f }
    { bitxor { } f }
} [
    first3 [
        [
            math-output-class/interval-special
        ] 2curry "output-classes" set-word-prop
    ] 2curry each-derived-op
] each

: twiddle-interval ( i1 -- i2 )
    dup [
        node get node-in-d
        [ value-class* integer class<= ] all?
        [ integral-closure ] when
    ] when ;

: (comparison-constraints) ( i1 i2 word class -- )
    node get [
        >r execute twiddle-interval 0 `input interval,
        r> 0 `output class,
    ] set-constraints ; inline

: comparison-constraints ( node true false -- )
    >r >r dup node set intervals dup [
        2dup
        r> \ f class-not (comparison-constraints)
        r> \ f (comparison-constraints)
    ] [
        r> r> 2drop 2drop
    ] if ; inline

{
    { < assume< assume>= }
    { <= assume<= assume> }
    { > assume> assume<= }
    { >= assume>= assume< }
} [
    first3 [
        [
            [ comparison-constraints ] with-scope
        ] 2curry "constraints" set-word-prop
    ] 2curry each-derived-op
] each

{
    alien-signed-1
    alien-unsigned-1
    alien-signed-2
    alien-unsigned-2
    alien-signed-4
    alien-unsigned-4
    alien-signed-8
    alien-unsigned-8
} [
    dup name>> {
        {
            [ "alien-signed-" ?head ]
            [ string>number 8 * 1- 2^ dup neg swap 1- [a,b] ]
        }
        {
            [ "alien-unsigned-" ?head ]
            [ string>number 8 * 2^ 1- 0 swap [a,b] ]
        }
    } cond 1array
    [ nip f swap ] curry "output-classes" set-word-prop
] each

! Associate intervals to classes
\ fixnum
most-negative-fixnum most-positive-fixnum [a,b]
"interval" set-word-prop

\ array-capacity
0 max-array-capacity [a,b]
"interval" set-word-prop

{
    { >fixnum fixnum }
    { >bignum bignum }
    { >float float }
} [
    [
        over node-in-d first value-interval*
        dup pick fits? [ drop f ] unless
        rot post-process
    ] curry "output-classes" set-word-prop
] assoc-each

! Removing overflow checks
: remove-overflow-check? ( #call -- ? )
    dup out-d>> first node-class
    [ fixnum class<= ] [ null eq? not ] bi and ;

{
    { + [ fixnum+fast ] }
    { +-integer-fixnum [ fixnum+fast ] }
    { - [ fixnum-fast ] }
    { * [ fixnum*fast ] }
    { *-integer-fixnum [ fixnum*fast ] }
    { shift [ fixnum-shift-fast ] }
    { fixnum+ [ fixnum+fast ] }
    { fixnum- [ fixnum-fast ] }
    { fixnum* [ fixnum*fast ] }
    { fixnum-shift [ fixnum-shift-fast ] }
} [
    [
        [ dup remove-overflow-check? ] ,
        [ f splice-quot ] curry ,
    ] { } make 1array define-optimizers
] assoc-each

! Remove redundant comparisons
: intervals-first2 ( #call -- first second )
    dup dup node-in-d first node-interval
    swap dup node-in-d second node-interval ;

: known-comparison? ( #call -- ? )
    intervals-first2 and ;

: perform-comparison ( #call word -- result )
    >r intervals-first2 r> execute ; inline

: foldable-comparison? ( #call word -- ? )
    >r dup known-comparison? [
        r> perform-comparison incomparable eq? not
    ] [
        r> 2drop f
    ] if ; inline

: fold-comparison ( #call word -- node )
    dupd perform-comparison 1array inline-literals ;

{
    { < interval< }
    { <= interval<= }
    { > interval> }
    { >= interval>= }
} [
    [
        [
            dup [ dupd foldable-comparison? ] curry ,
            [ fold-comparison ] curry ,
        ] { } make 1array define-optimizers
    ] curry each-derived-op
] assoc-each

! The following words are handled in a similar way except if
! the only consumer is a >fixnum we remove the overflow check
! too
: consumed-by? ( node word -- ? )
    swap sole-consumer
    dup #call? [ node-param eq? ] [ 2drop f ] if ;

: coerced-to-fixnum? ( #call -- ? )
    dup dup node-in-d [ node-class integer class<= ] with all?
    [ \ >fixnum consumed-by? ] [ drop f ] if ;

{
    { + [ [ >fixnum ] bi@ fixnum+fast ] }
    { - [ [ >fixnum ] bi@ fixnum-fast ] }
    { * [ [ >fixnum ] bi@ fixnum*fast ] }
} [
    >r derived-ops r> [
        [
            [
                dup remove-overflow-check?
                over coerced-to-fixnum? or
            ] ,
            [ f splice-quot ] curry ,
        ] { } make 1array define-optimizers
    ] curry each
] assoc-each

: convert-rem-to-and? ( #call -- ? )
    dup node-in-d {
        { [ 2dup first node-class integer class<= not ] [ f ] }
        { [ 2dup second node-literal integer? not ] [ f ] }
        { [ 2dup second node-literal power-of-2? not ] [ f ] }
        [ t ]
    } cond 2nip ;

: convert-mod-to-and? ( #call -- ? )
    dup dup node-in-d first node-interval 0 [a,inf] interval-subset?
    [ convert-rem-to-and? ] [ drop f ] if ;

: convert-mod-to-and ( #call -- node )
    dup
    dup in-d>> second node-literal 1-
    [ nip bitand ] curry f splice-quot ;

\ mod [
    {
        {
            [ dup convert-mod-to-and? ]
            [ convert-mod-to-and ]
        }
    } define-optimizers
] each-derived-op

\ rem {
    {
        [ dup convert-rem-to-and? ]
        [ convert-mod-to-and ]
    }
} define-optimizers

: fixnumify-bitand? ( #call -- ? )
    dup node-in-d second node-interval fixnum fits? ;

: fixnumify-bitand ( #call -- node )
    [ [ >fixnum ] bi@ fixnum-bitand ] f splice-quot ;

\ bitand {
    {
        [ dup fixnumify-bitand? ]
        [ fixnumify-bitand ]
    }
} define-optimizers

: convert-*-to-shift? ( #call -- ? )
    dup in-d>> second node-literal
    dup integer? [ power-of-2? ] [ drop f ] if ;

: convert-*-to-shift ( #call -- ? )
    dup dup in-d>> second node-literal log2
    [ nip fixnum-shift-fast ] curry
    f splice-quot ;

\ fixnum*fast {
    { [ dup convert-*-to-shift? ] [ convert-*-to-shift ] }
} define-optimizers

{ + - * / }
[ { number number } "input-classes" set-word-prop ] each

{ /f < > <= >= }
[ { real real } "input-classes" set-word-prop ] each

{ /i mod /mod }
[ { rational rational } "input-classes" set-word-prop ] each

{ bitand bitor bitxor bitnot shift }
[ { integer integer } "input-classes" set-word-prop ] each

{
    fcosh
    flog
    fsinh
    fexp
    fasin
    facosh
    fasinh
    ftanh
    fatanh
    facos
    fpow
    fatan
    fatan2
    fcos
    ftan
    fsin
    fsqrt
} [
    dup stack-effect
    [ in>> length real <repetition> "input-classes" set-word-prop ]
    [ out>> length float <repetition> "default-output-classes" set-word-prop ]
    2bi
] each