extra/trees/splay/splay.factor

   1 ! Copyright (c) 2005 Mackenzie Straight.
   2 ! See http://factorcode.org/license.txt for BSD license.
   3 USING: arrays kernel math namespaces sequences assocs parser
   4 prettyprint.backend trees generic math.order ;
   5 IN: trees.splay
   6
   7 TUPLE: splay < tree ;
   8
   9 : <splay> ( -- tree )
  10     \ splay new-tree ;
  11
  12 : rotate-right ( node -- node )
  13     dup node-left
  14     [ node-right swap set-node-left ] 2keep
  15     [ set-node-right ] keep ;
  16
  17 : rotate-left ( node -- node )
  18     dup node-right
  19     [ node-left swap set-node-right ] 2keep
  20     [ set-node-left ] keep ;
  21
  22 : link-right ( left right key node -- left right key node )
  23     swap >r [ swap set-node-left ] 2keep
  24     nip dup node-left r> swap ;
  25
  26 : link-left ( left right key node -- left right key node )
  27     swap >r rot [ set-node-right ] 2keep
  28     drop dup node-right swapd r> swap ;
  29
  30 : cmp ( key node -- obj node -1/0/1 )
  31     2dup node-key key-side ;
  32
  33 : lcmp ( key node -- obj node -1/0/1 )
  34     2dup node-left node-key key-side ;
  35
  36 : rcmp ( key node -- obj node -1/0/1 )
  37     2dup node-right node-key key-side ;
  38
  39 DEFER: (splay)
  40
  41 : splay-left ( left right key node -- left right key node )
  42     dup node-left [
  43         lcmp 0 < [ rotate-right ] when
  44         dup node-left [ link-right (splay) ] when
  45     ] when ;
  46
  47 : splay-right ( left right key node -- left right key node )
  48     dup node-right [
  49         rcmp 0 > [ rotate-left ] when
  50         dup node-right [ link-left (splay) ] when
  51     ] when ;
  52
  53 : (splay) ( left right key node -- left right key node )
  54     cmp dup 0 <
  55     [ drop splay-left ] [ 0 > [ splay-right ] when ] if ;
  56
  57 : assemble ( head left right node -- root )
  58     [ node-right swap set-node-left ] keep
  59     [ node-left swap set-node-right ] keep
  60     [ swap node-left swap set-node-right ] 2keep
  61     [ swap node-right swap set-node-left ] keep ;
  62
  63 : splay-at ( key node -- node )
  64     >r >r T{ node } clone dup dup r> r>
  65     (splay) nip assemble ;
  66
  67 : splay ( key tree -- )
  68     [ tree-root splay-at ] keep set-tree-root ;
  69
  70 : splay-split ( key tree -- node node )
  71     2dup splay tree-root cmp 0 < [
  72         nip dup node-left swap f over set-node-left
  73     ] [
  74         nip dup node-right swap f over set-node-right swap
  75     ] if ;
  76
  77 : get-splay ( key tree -- node ? )
  78     2dup splay tree-root cmp 0 = [
  79         nip t
  80     ] [
  81         2drop f f
  82     ] if ;
  83
  84 : get-largest ( node -- node )
  85     dup [ dup node-right [ nip get-largest ] when* ] when ;
  86
  87 : splay-largest ( node -- node )
  88     dup [ dup get-largest node-key swap splay-at ] when ;
  89
  90 : splay-join ( n2 n1 -- node )
  91     splay-largest [
  92         [ set-node-right ] keep
  93     ] [
  94         drop f
  95     ] if* ;
  96
  97 : remove-splay ( key tree -- )
  98     tuck get-splay nip [
  99         dup dec-count
 100         dup node-right swap node-left splay-join
 101         swap set-tree-root
 102     ] [ drop ] if* ;
 103
 104 : set-splay ( value key tree -- )
 105     2dup get-splay [ 2nip set-node-value ] [
 106        drop dup inc-count
 107        2dup splay-split rot
 108        >r >r swapd r> node boa r> set-tree-root
 109     ] if ;
 110
 111 : new-root ( value key tree -- )
 112     [ 1 swap set-tree-count ] keep
 113     >r swap <node> r> set-tree-root ;
 114
 115 M: splay set-at ( value key tree -- )
 116     dup tree-root [ set-splay ] [ new-root ] if ;
 117
 118 M: splay at* ( key tree -- value ? )
 119     dup tree-root [
 120         get-splay >r dup [ node-value ] when r>
 121     ] [
 122         2drop f f
 123     ] if ;
 124
 125 M: splay delete-at ( key tree -- )
 126     dup tree-root [ remove-splay ] [ 2drop ] if ;
 127
 128 M: splay new-assoc
 129     2drop <splay> ;
 130
 131 : >splay ( assoc -- tree )
 132     T{ splay f f 0 } assoc-clone-like ;
 133
 134 : SPLAY{
 135     \ } [ >splay ] parse-literal ; parsing
 136
 137 M: splay assoc-like
 138     drop dup splay? [ >splay ] unless ;
 139
 140 M: splay pprint-delims drop \ SPLAY{ \ } ;