! Copyright (C) 2006, 2008 Slava Pestov.
! See http://factorcode.org/license.txt for BSD license.
USING: classes io kernel kernel.private math.parser namespaces
optimizer prettyprint prettyprint.backend sequences words arrays
match macros assocs sequences.private generic combinators
sorting math quotations accessors inference inference.backend
inference.dataflow optimizer.specializers generator ;
IN: optimizer.debugger

! A simple tool for turning dataflow IR into quotations, for
! debugging purposes.

GENERIC: node>quot ( ? node -- )

TUPLE: comment node text ;

M: comment pprint*
    "( " over comment-text " )" 3append
    swap comment-node present-text ;

: comment, ( ? node text -- )
    rot [ \ comment boa , ] [ 2drop ] if ;

: values% ( prefix values -- )
    swap [
        %
        dup value? [
            value-literal unparse %
        ] [
            "@" % unparse %
        ] if
    ] curry each ;

: effect-str ( node -- str )
    [
        " " over in-d>> values%
        " r: " over in-r>> values%
        " --" %
        " " over out-d>> values%
        " r: " swap out-r>> values%
    ] "" make rest ;

MACRO: match-choose ( alist -- )
    [ [ ] curry ] assoc-map [ match-cond ] curry ;

MATCH-VARS: ?a ?b ?c ;

: pretty-shuffle ( in out -- word/f )
    2array {
        { { { ?a } { ?a } } [ ] }
        { { { ?a ?b } { ?a ?b } } [ ] }
        { { { ?a ?b ?c } { ?a ?b ?c } } [ ] }
        { { { ?a } { } } [ drop ] }
        { { { ?a ?b } { } } [ 2drop ] }
        { { { ?a ?b ?c } { } } [ 3drop ] }
        { { { ?a } { ?a ?a } } [ dup ] }
        { { { ?a ?b } { ?a ?b ?a ?b } } [ 2dup ] }
        { { { ?a ?b ?c } { ?a ?b ?c ?a ?b ?c } } [ 3dup ] }
        { { { ?a ?b } { ?a ?b ?a } } [ over ] }
        { { { ?b ?a } { ?a ?b } } [ swap ] }
        { { { ?a ?b } { ?b ?a ?b } } [ tuck ] }
        { { { ?a ?b ?c } { ?a ?b ?c ?a } } [ pick ] }
        { { { ?a ?b ?c } { ?c ?a ?b } } [ -rot ] }
        { { { ?a ?b ?c } { ?b ?c ?a } } [ rot ] }
        { { { ?a ?b } { ?b } } [ nip ] }
        { { { ?a ?b ?c } { ?c } } [ 2nip ] }
        { _ f }
    } match-choose ;

M: #shuffle node>quot
    dup [ in-d>> ] [ out-d>> ] bi pretty-shuffle
    [ % ] [ >r drop t r> ] if*
    dup effect-str "#shuffle: " prepend comment, ;

: pushed-literals ( node -- seq )
    out-d>> [ value-literal literalize ] map ;

M: #push node>quot nip pushed-literals % ;

DEFER: dataflow>quot

: #call>quot ( ? node -- )
    dup param>> dup ,
    [ dup effect-str ] [ "empty call" ] if comment, ;

M: #call node>quot #call>quot ;

M: #call-label node>quot #call>quot ;

M: #label node>quot
    [
        dup param>> literalize ,
        dup #label-loop? "#loop: " "#label: " ?
        over param>> name>> append comment,
    ] 2keep
    node-child swap dataflow>quot , \ call ,  ;

M: #if node>quot
    [ "#if" comment, ] 2keep
    children>> swap [ dataflow>quot ] curry map %
    \ if , ;

M: #dispatch node>quot
    [ "#dispatch" comment, ] 2keep
    children>> swap [ dataflow>quot ] curry map ,
    \ dispatch , ;

M: #>r node>quot nip in-d>> length \ >r <array> % ;

M: #r> node>quot nip out-d>> length \ r> <array> % ;

M: object node>quot
    [
        dup class name>> %
        " " %
        dup param>> unparse %
        " " %
        dup effect-str %
    ] "" make comment, ;

: (dataflow>quot) ( ? node -- )
    dup [
        2dup node>quot successor>> (dataflow>quot)
    ] [
        2drop
    ] if ;

: dataflow>quot ( node ? -- quot )
    [ swap (dataflow>quot) ] [ ] make ;

: optimized-quot. ( quot ? -- )
    #! Print dataflow IR for a quotation. Flag indicates if
    #! annotations should be printed or not.
    >r dataflow optimize r> dataflow>quot pprint nl ;

: optimized-word. ( word ? -- ) >r specialized-def r> optimized-quot. ;

SYMBOL: pass-count
SYMBOL: words-called
SYMBOL: generics-called
SYMBOL: methods-called
SYMBOL: intrinsics-called
SYMBOL: node-count

: count-optimization-passes ( node n -- node n )
    >r optimize-1
    [ r> 1+ count-optimization-passes ] [ r> ] if ;

: make-report ( word -- assoc )
    [
        word-dataflow nip 1 count-optimization-passes pass-count set

        H{ } clone words-called set
        H{ } clone generics-called set
        H{ } clone methods-called set
        H{ } clone intrinsics-called set

        0 swap [
            >r 1+ r>
            dup #call? [
                param>> {
                    { [ dup "intrinsics" word-prop over "if-intrinsics" word-prop or ] [ intrinsics-called ] }
                    { [ dup generic? ] [ generics-called ] }
                    { [ dup method-body? ] [ methods-called ] }
                    [ words-called ]
                } cond 1 -rot get at+
            ] [
                drop
            ] if
        ] each-node
        node-count set
    ] H{ } make-assoc ;

: report. ( report -- )
    [
        "==== Optimization passes:" print
        pass-count get .
        nl

        "==== Total number of dataflow nodes:" print
        node-count get .

        {
            { generics-called "==== Generic word calls:" }
            { words-called "==== Ordinary word calls:" }
            { methods-called "==== Non-inlined method calls:" }
            { intrinsics-called "==== Open-coded intrinsic calls:" }
        } [
            nl print get keys natural-sort stack.
        ] assoc-each
    ] bind ;

: optimizer-report. ( word -- )
    make-report report. ;