roc/examples/benchmarks/RBTreeInsert.roc
2021-01-19 20:22:08 +01:00

107 lines
2.9 KiB
Plaintext

app "rbtree-insert"
packages { base: "platform" }
imports [base.Task]
provides [ main ] to base
main : Task.Task {} []
main =
tree : RedBlackTree I64 {}
tree = insert 0 {} Empty
tree
|> show
|> Task.putLine
show : RedBlackTree I64 {} -> Str
show = \tree -> showRBTree tree Str.fromInt (\{} -> "{}")
showRBTree : RedBlackTree k v, (k -> Str), (v -> Str) -> Str
showRBTree = \tree, showKey, showValue ->
when tree is
Empty -> "Empty"
Node color key value left right ->
sColor = showColor color
sKey = showKey key
sValue = showValue value
sL = nodeInParens left showKey showValue
sR = nodeInParens right showKey showValue
"Node \(sColor) \(sKey) \(sValue) \(sL) \(sR)"
nodeInParens : RedBlackTree k v, (k -> Str), (v -> Str) -> Str
nodeInParens = \tree, showKey, showValue ->
when tree is
Empty -> showRBTree tree showKey showValue
Node _ _ _ _ _ ->
inner = showRBTree tree showKey showValue
"(\(inner))"
showColor : NodeColor -> Str
showColor = \color ->
when color is
Red -> "Red"
Black -> "Black"
NodeColor : [ Red, Black ]
RedBlackTree k v : [ Node NodeColor k v (RedBlackTree k v) (RedBlackTree k v), Empty ]
Key k : Num k
insert : Key k, v, RedBlackTree (Key k) v -> RedBlackTree (Key k) v
insert = \key, value, dict ->
when insertHelp key value dict is
Node Red k v l r ->
Node Black k v l r
x ->
x
insertHelp : (Key k), v, RedBlackTree (Key k) v -> RedBlackTree (Key k) v
insertHelp = \key, value, dict ->
when dict is
Empty ->
# New nodes are always red. If it violates the rules, it will be fixed
# when balancing.
Node Red key value Empty Empty
Node nColor nKey nValue nLeft nRight ->
when Num.compare key nKey is
LT ->
balance nColor nKey nValue (insertHelp key value nLeft) nRight
EQ ->
Node nColor nKey value nLeft nRight
GT ->
balance nColor nKey nValue nLeft (insertHelp key value nRight)
balance : NodeColor, k, v, RedBlackTree k v, RedBlackTree k v -> RedBlackTree k v
balance = \color, key, value, left, right ->
when right is
Node Red rK rV rLeft rRight ->
when left is
Node Red lK lV lLeft lRight ->
Node
Red
key
value
(Node Black lK lV lLeft lRight)
(Node Black rK rV rLeft rRight)
_ ->
Node color rK rV (Node Red key value left rLeft) rRight
_ ->
when left is
Node Red lK lV (Node Red llK llV llLeft llRight) lRight ->
Node
Red
lK
lV
(Node Black llK llV llLeft llRight)
(Node Black key value lRight right)
_ ->
Node color key value left right