Paste: tries sketch

USING: accessors arrays assocs combinators combinators.smart fry
inverse kernel locals math sequences strings ;
IN: tries

TUPLE: trie value { children assoc } ;
: <trie> ( -- trie ) f { } trie boa ;

: trie-find ( str trie -- value/f ? )
    over empty?
    [ value>> nip t ]
    [ [ unclip 1string ]
      [ children>> ] bi* at dup
      [ trie-find ]
      [ nip f ] if
    ] if ; inline recursive

:: trie-insert ( trie key value  -- )
    key empty?
    [ trie value >>value drop ]
    [ key unclip 1string :> ( ks k )
      trie children>> k swap at dup
      [ ks value trie-insert ]  ! branch already present
      [ drop
        <trie> :> branch
        trie [ { k branch } suffix ] change-children drop
        branch ks value trie-insert ] if 
    ] if ; inline recursive

: trie-count ( trie -- x )
    [ value>> [ 1 ] [ 0 ] if ]
    [ children>> [ second trie-count ] map sum ] bi + ; inline recursive

DEFER: collect

: trie-collapse ( tr -- assoc )
    dup children>> empty? 
    [ value>> "" swap 2array 1array ] 
    [ [ children>> [ collect ] { } assoc>map concat
    ] keep value>> [ "" swap 2array suffix ] when* 
    ] if ; inline recursive

: collect ( str trie -- assoc )
    trie-collapse swap '[ [ _ prepend ] dip ] assoc-map
    ; inline recursive

! want fried sequences!!!
M: trie at* trie-find ;
M: trie assoc-size trie-count ;
M: trie >alist trie-collapse ;

INSTANCE: trie assoc

M: trie set-at trie-insert ;
: >trie ( assoc -- trie ) <trie> [ '[ _ -rot trie-insert ] assoc-each ] keep ;
M: trie assoc-like drop dup trie? [ >trie ] unless ;