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