1	# Shape invariants depend on tree structure, never on the stored element type.
2	# Native capacity is a U32; extent is unbounded and does not assert no overflow.
3	import Base
4	import ./storage_tree_model.bend as Tree
5	
6	type Both<-Left: Data,-Right: Data> is Data:
7	  Both{left: Left,right: Right}
8	
9	def first(-Left: Data,-Right: Data,evidence: Both<Left,Right>) -> Left:
10	  Both{left,right} = evidence
11	  left
12	
13	def second(-Left: Data,-Right: Data,evidence: Both<Left,Right>) -> Right:
14	  Both{left,right} = evidence
15	  right
16	
17	law Balanced:
18	  for -Element: Data
19	  for depth: Nat
20	  for tree: Tree.Tree<Element>
21	  Data
22	
23	def Balanced(Element,depth,tree):
24	  match depth tree:
25	    case 0n Tree.Leaf{value}: Unit
26	    case 0n Tree.Node{left,right}: Empty
27	    case 1n+rest Tree.Leaf{value}: Empty
28	    case 1n+ +rest Tree.Node{left,right}:
29	      Both<Balanced(Element,rest,left),Balanced(Element,rest,right)>
30	
31	def extent(depth: Nat) -> Nat:
32	  match depth:
33	    case 0n: 1n
34	    case 1n+rest: Nat.double(extent(rest))
35	
36	def fresh(-Element: Data,+depth: Nat,+value: Element) -> Balanced(Element,depth,Tree.fresh(Element,depth,value)):
37	  match depth:
38	    case 0n: Unit{}
39	    case 1n+ +rest: Both{fresh(Element,rest,value),fresh(Element,rest,value)}
40	
41	def store_node(-Element: Data,+depth: Nat,+left: Tree.Tree<Element>,+right: Tree.Tree<Element>,
42	  +half: U32,+index: U32,+value: Element,+choose_left: Bool,
43	  left_balanced: Balanced(Element,depth,left),right_balanced: Balanced(Element,depth,right),
44	  left_stored: Balanced(Element,depth,Tree.store(Element,left,half,index,value)),
45	  right_stored: Balanced(Element,depth,Tree.store(Element,right,half,(index - half : U32),value))) ->
46	  Balanced(Element,1n+depth,Bool.pick(Tree.Tree<Element>,choose_left,
47	    Tree.Node{Tree.store(Element,left,half,index,value),right},
48	    Tree.Node{left,Tree.store(Element,right,half,(index - half : U32),value)})):
49	  match choose_left:
50	    case True{}: Both{left_stored,right_balanced}
51	    case False{}: Both{left_balanced,right_stored}
52	
53	def store(-Element: Data,+depth: Nat,+tree: Tree.Tree<Element>,+capacity: U32,+index: U32,+value: Element,
54	  balanced: Balanced(Element,depth,tree)) -> Balanced(Element,depth,Tree.store(Element,tree,capacity,index,value)):
55	  match depth tree:
56	    case 0n Tree.Leaf{previous}: Unit{}
57	    case 0n Tree.Node{left,right}: match balanced:
58	    case 1n+rest Tree.Leaf{previous}: match balanced:
59	    case 1n+ +rest Tree.Node{+left,+right}:
60	      Both{+left_balanced,+right_balanced} = balanced
61	      store_node(Element,rest,left,right,U32.shr(capacity),index,value,U32.is_lt(index,U32.shr(capacity)),
62	        left_balanced,right_balanced,
63	        store(Element,rest,left,U32.shr(capacity),index,value,left_balanced),
64	        store(Element,rest,right,U32.shr(capacity),(index - U32.shr(capacity) : U32),value,right_balanced))
65	
66	def write(-Element: Data,+depth: Nat,+tree: Tree.Tree<Element>,+index: U32,+value: Element,
67	  balanced: Balanced(Element,depth,tree)) -> Balanced(Element,depth,Tree.write(Element,tree,index,value)):
68	  store(Element,depth,tree,Tree.capacity(Element,tree),U32.and(index,U32.sub(Tree.capacity(Element,tree),1)),value,balanced)
69	
70	def capacity(-Element: Data,+depth: Nat,+tree: Tree.Tree<Element>,balanced: Balanced(Element,depth,tree)) ->
71	  {Tree.capacity(Element,tree) == U32.shln(1,depth) : U32}:
72	  match depth tree:
73	    case 0n Tree.Leaf{value}: {==}
74	    case 0n Tree.Node{left,right}: match balanced:
75	    case 1n+rest Tree.Leaf{value}: match balanced:
76	    case 1n+ +rest Tree.Node{+left,+right}:
77	      Both{left_balanced,right_balanced} = balanced
78	      %capacity(Element,rest,left,left_balanced) : {U32.shl(Tree.capacity(Element,left)) == U32.shl(_) : U32}
79	      {==}
80	
81	def unique_depth(-Element: Data,tree: Tree.Tree<Element>,left_depth: Nat,right_depth: Nat,
82	  left: Balanced(Element,left_depth,tree),right: Balanced(Element,right_depth,tree)) -> {left_depth == right_depth : Nat}:
83	  match tree left_depth right_depth:
84	    case Tree.Leaf{value} 0n 0n: {==}
85	    case Tree.Leaf{value} 0n 1n+right_depth: match right:
86	    case Tree.Leaf{value} 1n+left_depth right_depth: match left:
87	    case Tree.Node{left_tree,right_tree} 0n right_depth: match left:
88	    case Tree.Node{left_tree,right_tree} 1n+left_depth 0n: match right:
89	    case Tree.Node{+left_tree,right_tree} 1n+ +left_depth 1n+ +right_depth:
90	      Both{left_balanced,left_other} = left
91	      Both{right_balanced,right_other} = right
92	      Equal.cong(Nat,Nat,depth => 1n+depth,left_depth,right_depth,unique_depth(Element,left_tree,left_depth,right_depth,left_balanced,right_balanced))
93	
94	# Balanced trees at the same depth expose the same native capacity.
95	def same_capacity(-Element: Data,+depth: Nat,+left: Tree.Tree<Element>,+right: Tree.Tree<Element>,
96	  left_balanced: Balanced(Element,depth,left),right_balanced: Balanced(Element,depth,right)) ->
97	  {Tree.capacity(Element,left) == Tree.capacity(Element,right) : U32}:
98	  Equal.trans(U32,Tree.capacity(Element,left),U32.shln(1,depth),Tree.capacity(Element,right),capacity(Element,depth,left,left_balanced),
99	    Equal.sym(U32,Tree.capacity(Element,right),U32.shln(1,depth),capacity(Element,depth,right,right_balanced)))
100	
101	def transfer_room(-Element: Data,+depth: Nat,+left: Tree.Tree<Element>,+right: Tree.Tree<Element>,+count: Nat,
102	  left_balanced: Balanced(Element,depth,left),right_balanced: Balanced(Element,depth,right),room: {Nat.is_le(count,U32.to_nat(Tree.capacity(Element,left))) == True{} : Bool}) ->
103	  {Nat.is_le(count,U32.to_nat(Tree.capacity(Element,right))) == True{} : Bool}:
104	  %same_capacity(Element,depth,left,right,left_balanced,right_balanced) : {Nat.is_le(count,U32.to_nat(_)) == True{} : Bool}
105	  room