USING: accessors arrays assocs combinators combinators.smart fry inverse kernel locals math sequences strings ; IN: tries TUPLE: trie value { children assoc } ; : ( -- 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 :> 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 ) [ '[ _ -rot trie-insert ] assoc-each ] keep ; M: trie assoc-like drop dup trie? [ >trie ] unless ;