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

{ + bignum+ float+ fixnum+fast } {
    { { number 0 } [ drop ] }
    { { 0 number } [ nip ] }
} define-identities

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

{ - fixnum- bignum- float- fixnum-fast } {
    { { number 0 } [ drop ] }
    { { @ @ } [ 2drop 0 ] }
} define-identities

{ < fixnum< bignum< float< } {
    { { @ @ } [ 2drop f ] }
} define-identities

{ <= fixnum<= bignum<= float<= } {
    { { @ @ } [ 2drop t ] }
} define-identities

{ > fixnum> bignum> float>= } {
    { { @ @ } [ 2drop f ] }
} define-identities

{ >= fixnum>= bignum>= float>= } {
    { { @ @ } [ 2drop t ] }
} define-identities

{ * fixnum* bignum* float* } {
    { { number 1 } [ drop ] }
    { { 1 number } [ nip ] }
    { { number 0 } [ nip ] }
    { { 0 number } [ drop ] }
    { { number -1 } [ drop 0 swap - ] }
    { { -1 number } [ nip 0 swap - ] }
} define-identities

{ / fixnum/i bignum/i float/f } {
    { { number 1 } [ drop ] }
    { { number -1 } [ drop 0 swap - ] }
} define-identities

{ fixnum-mod bignum-mod } {
    { { number 1 } [ 2drop 0 ] }
} define-identities

{ bitand fixnum-bitand bignum-bitand } {
    { { number -1 } [ drop ] }
    { { -1 number } [ nip ] }
    { { @ @ } [ drop ] }
    { { number 0 } [ nip ] }
    { { 0 number } [ drop ] }
} define-identities

{ bitor fixnum-bitor bignum-bitor } {
    { { number 0 } [ drop ] }
    { { 0 number } [ nip ] }
    { { @ @ } [ drop ] }
    { { number -1 } [ nip ] }
    { { -1 number } [ drop ] }
} define-identities

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

{ shift fixnum-shift fixnum-shift-fast bignum-shift } {
    { { 0 number } [ drop ] }
    { { number 0 } [ drop ] }
} define-identities

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

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

: math-output-class ( node min -- newclass )
    #! if min is f, it means we just want to use the declared
    #! output class from the "infer-effect".
    dup [
        swap node-in-d
        [ value-class* math-closure math-class-max ] each
    ] [
        2drop f
    ] if ;

: 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 ] 2apply ;

: 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 min word -- classes intervals )
    pick >r
    >r over r>
    math-output-interval-1
    >r math-output-class r>
    r> post-process ; inline

{
    { 1+ integer interval-1+ }
    { 1- integer interval-1- }
    { neg integer interval-neg }
    { shift integer interval-recip }
    { bitnot fixnum interval-bitnot }
    { fixnum-bitnot f interval-bitnot }
    { bignum-bitnot f interval-bitnot }
    { 2/ fixnum interval-2/ }
    { sq integer f }
} [
    first3 [
        math-output-class/interval-1
    ] 2curry "output-classes" set-word-prop
] each

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

: 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 min word -- classes intervals )
    pick >r
    >r over r>
    math-output-interval-2
    >r math-output-class r>
    r> post-process ; inline

{
    { + integer interval+ }
    { - integer interval- }
    { * integer interval* }
    { / rational interval/ }
    { /i integer interval/i }

    { fixnum+ f interval+ }
    { fixnum+fast f interval+ }
    { fixnum- f interval- }
    { fixnum-fast f interval- }
    { fixnum* f interval* }
    { fixnum*fast f interval* }
    { fixnum/i f interval/i }

    { bignum+ f interval+ }
    { bignum- f interval- }
    { bignum* f interval* }
    { bignum/i f interval/i }
    { bignum-shift f interval-shift-safe }

    { float+ f interval+ }
    { float- f interval- }
    { float* f interval* }
    { float/f f interval/ }

    { min fixnum interval-min }
    { max fixnum interval-max }
} [
    first3 [
        math-output-class/interval-2
    ] 2curry "output-classes" set-word-prop
] each

{ fixnum-shift fixnum-shift-fast shift } [
    [
        dup
        node-in-d second value-interval*
        -1./0. 0 [a,b] interval-subset? fixnum integer ?
        \ interval-shift-safe
        math-output-class/interval-2
    ] "output-classes" set-word-prop
] 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 fixnum mod-range }
    { fixnum-mod f mod-range }
    { bignum-mod f mod-range }
    { float-mod f mod-range }

    { rem integer rem-range }

    { bitand fixnum bitand-range }
    { fixnum-bitand f bitand-range }

    { bitor fixnum f }
    { bitxor fixnum f }
} [
    first3 [
        math-output-class/interval-special
    ] 2curry "output-classes" set-word-prop
] 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> general-t (comparison-constraints)
        r> \ f (comparison-constraints)
    ] [
        r> r> 2drop 2drop
    ] if ; inline

{
    { < assume< assume>= }
    { <= assume<= assume> }
    { > assume> assume<= }
    { >= assume>= assume< }

    { fixnum< assume< assume>= }
    { fixnum<= assume<= assume> }
    { fixnum> assume> assume<= }
    { fixnum>= assume>= assume< }

    { bignum< assume< assume>= }
    { bignum<= assume<= assume> }
    { bignum> assume> assume<= }
    { bignum>= assume>= assume< }

    { float< assume< assume>= }
    { float<= assume<= assume> }
    { float> assume> assume<= }
    { float>= assume>= assume< }
} [
    first3
    [
        [ comparison-constraints ] with-scope
    ] 2curry "constraints" set-word-prop
] 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 word-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 node-out-d first node-class fixnum class< ;

{
    { + [ fixnum+fast ] }
    { - [ fixnum-fast ] }
    { * [ fixnum*fast ] }
    { fixnum+ [ fixnum+fast ] }
    { fixnum- [ fixnum-fast ] }
    { fixnum* [ fixnum*fast ] }
    ! these are here as an optimization. if they weren't given
    ! explicitly, the same would be inferred after an extra
    ! optimization step (see optimistic-inline?)
    { 1+ [ 1 fixnum+fast ] }
    { 1- [ 1 fixnum-fast ] }
    { 2/ [ -1 fixnum-shift ] }
    { neg [ 0 swap fixnum-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>= }

    { fixnum< interval< }
    { fixnum<= interval<= }
    { fixnum> interval> }
    { fixnum>= interval>= }

    { bignum< interval< }
    { bignum<= interval<= }
    { bignum> interval> }
    { bignum>= interval>= }

    { float< interval< }
    { float<= interval<= }
    { float> interval> }
    { float>= interval>= }
} [
    [
        dup [ dupd foldable-comparison? ] curry ,
        [ fold-comparison ] curry ,
    ] { } make 1array define-optimizers
] 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 ;

: coereced-to-fixnum? ( #call -- ? )
    \ >fixnum consumed-by? ;

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

: fixnum-shift-fast-pos? ( node -- ? )
    #! Shifting 1 to the left won't overflow if the shift
    #! count is small enough
    dup dup node-in-d first node-literal 1 = [
        dup node-in-d second node-interval
        0 cell-bits tag-bits get - 2 - [a,b] interval-subset?
    ] [ drop f ] if ;

: fixnum-shift-fast-neg? ( node -- ? )
    #! Shifting any number to the right won't overflow if the
    #! shift count is small enough
    dup node-in-d second node-interval
    cell-bits 1- neg 0 [a,b] interval-subset? ;

: fixnum-shift-fast? ( node -- ? )
    dup fixnum-shift-fast-pos?
    [ drop t ] [ fixnum-shift-fast-neg? ] if ;

\ fixnum-shift {
    {
        [ dup fixnum-shift-fast? ]
        [ [ fixnum-shift-fast ] f splice-quot ]
    }
} define-optimizers