Paste: object-graph

USING: accessors arrays assocs classes combinators deques
dlists fry images.loader images.viewer io io.encodings.utf8
io.files io.files.unique io.launcher kernel locals make math
math.functions math.parser mirrors namespaces sequences strings
;
IN: object-graph

SYMBOLS: node-counter node-hash edges unprocessed-nodes ;

: init ( -- )
    0 node-counter set
    H{ } clone node-hash set
    V{ } clone edges set
    <dlist> unprocessed-nodes set ;

: new-node-counter ( -- n )
    node-counter [ dup 1 + ] change ;

TUPLE: literal string ;
C: <literal> literal

TUPLE: node-reference object ;
C: <node-reference> node-reference

TUPLE: slot-reference slot-name object ;
C: <slot-reference> slot-reference

GENERIC: add-edge ( from to -- )

M: literal add-edge
    2drop ;
M: node-reference add-edge
    2array edges get push ;

: add-node ( object -- )
    [ new-node-counter ] dip node-hash get set-at ;

: escape-string ( string -- string' )
    [
        [
            {
                { CHAR: \n [ "\\\n" % ] }
                { CHAR: \t [ "\\\t" % ] }
                { CHAR: \\ [ "\\\\" % ] }
                { CHAR: \" [ "\\\"" % ] }
                { CHAR: \r [ "\\\r" % ] }
                [ , ]
            } case
        ] each
    ] "" make ;

GENERIC: traverse ( obj -- result )

M: number traverse
    number>string <literal> ;
M: complex traverse
    >rect [ number>string ] bi@ "+" glue "i" append <literal> ;
M: string traverse
    escape-string "\\\"" dup surround <literal> ;
M: f traverse
    drop "f" <literal> ;
M: object traverse
    dup node-hash get key? [
        [ unprocessed-nodes get push-back ] [ add-node ] [ ] tri
    ] unless <node-reference> ;

GENERIC: output-target ( target -- )

M: literal output-target
    string>> write ;
M: node-reference output-target
    drop "Â·" write ;

: output-field ( name value -- )
    " | { " write
    [ drop write " | " write ]
    [ [ "<" "> " surround write ] dip output-target ] 2bi
    " } " write ;

: add-edges ( object -- )
    [ <mirror> ] keep '[
        traverse
        [ [ _ <slot-reference> ] dip add-edge ]
        [ output-field ] 2bi
    ] assoc-each ;

: output-header ( object -- )
    [ node-hash get at number>string write ] keep
    " [label=\"{ " write
    class name>> write ;

GENERIC: output-object ( object -- )

M: object output-object
    [ output-header ] [ add-edges ] bi "}\"]" print ;

: add-nths ( seq -- )
    dup '[
        [ traverse ] [ number>string ] bi*
        [ _ <slot-reference> swap add-edge ]
        [ " | <" "> " surround write output-target ] 2bi
    ] each-index ;

M: sequence output-object
    [ output-header ] [ add-nths ] bi "}\"]" print ;

GENERIC# add-key-value 2 ( elt assoc prefix -- )

M: literal add-key-value
    2drop string>> write ;

M:: node-reference add-key-value ( elt assoc prefix -- )
    elt object>> :> elt'
    elt' node-hash get at :> id
    prefix id number>string append :> slot-name
    slot-name assoc <slot-reference> elt add-edge
    "<" slot-name "> . " 3append write ;

: add-key-values ( assoc -- )
    dup '[
        " | { " write
        [ traverse ] bi@
        _ [ "k" add-key-value " | " write ]
        [ "v" add-key-value ] bi-curry bi*
        " } " write
    ] assoc-each ;

M: assoc output-object
    [ output-header ] [ add-key-values ] bi "}\"]" print ;

: build-object-graph ( object -- )
    traverse
    dup literal?
    [ string>> "0 [label=\"" "\"]" surround print ]
    [ drop unprocessed-nodes get [ output-object ] slurp-deque ] if ;

: output-slot-reference ( slot-reference -- )
    [ object>> node-hash get at number>string ]
    [ slot-name>> ] bi ":" glue write ;

: output-edge ( from to -- )
    [ output-slot-reference ] dip
    " -> " write
    object>> node-hash get at number>string print ;

: output-edges ( -- )
    edges get [ output-edge ] assoc-each ;

: output-object-graph ( object -- )
    [
        init
        "digraph structs {" print
        "node [shape=record];" print
        build-object-graph
        output-edges
        "}" print
    ] with-scope ;

: object-graph>png ( object -- filename )
    "object-graph" "dot" [
        [ utf8 [ output-object-graph ] with-file-writer ]
        [ { "dot" "-Tpng" "-O" } swap suffix try-process ]
        [ ".png" append ] tri
    ] cleanup-unique-file ;

: object-graph. ( object -- )
    object-graph>png load-image image. ;