# MIT License, Copyright (c) 2024 Marvin Borner
# finger tree implementation, great for sequences / searches
# originally by R. Hinze, R. Paterson "Finger Trees: a simple general-purpose data structure"
# efficient translation to LC by me

:import std/Combinator .
:import std/List L
:import std/Number N

# === Node ===
# Scott-style tagged union, 2 tags

# (Node a) = (Node 2 a a) | (Node3 a a a)

# tagged two elements (tag 0)
node2 [[[[0 3 2]]]] ⧗ a → a → (Node a)

# tagged three elements (tag 1)
node3 [[[[[1 4 3 2]]]]] ⧗ a → a → a → (Node a)

foldr-node [[[0 case-node3 case-node2]]] ⧗ (a → b → b) → b → (Node a) → b
	case-node2 [[4 1 (4 0 3)]]
	case-node3 [[[5 2 (5 1 (5 0 4))]]]

:test (foldr-node L.cons L.empty (node2 'a' 'b')) ("ab")
:test (foldr-node L.cons L.empty (node3 'a' 'b' 'c')) ("abc")

foldl-node [[[0 case-node3 case-node2]]] ⧗ (b → a → b) → b → (Node a) → b
	case-node2 [[4 (4 3 1) 0]]
	case-node3 [[[5 (5 (5 4 2) 1) 0]]]

:test (foldl-node \L.cons L.empty (node2 'a' 'b')) ("ba")
:test (foldl-node \L.cons L.empty (node3 'a' 'b' 'c')) ("cba")

# === Digit ===
# Scott-style tagged union with constant cons/snoc, 4 tags
# alternative would be Parigot lists, but size check requires y-rec

# tagged single element (tag 0)
one [[[[[0 4]]]]] ⧗ a → (Digit a)

{}‣ one

# tagged two elements (tag 1)
two [[[[[[1 5 4]]]]]] ⧗ a → a → (Digit a)

# tagged three elements (tag 2)
three [[[[[[[2 6 5 4]]]]]]] ⧗ a → a → a → (Digit a)

# tagged four elements (tag 3)
four [[[[[[[[3 7 6 5 4]]]]]]]] ⧗ a → a → a → a → (Digit a)

four? [0 [[[[k]]]] [[[ki]]] [[ki]] [ki]] ⧗ (Digit a) → Boolean

:test (four? (one i)) (ki)
:test (four? (two i i)) (ki)
:test (four? (three i i i)) (ki)
:test (four? (four i i i i)) (k)

# returns first element of digit
digit-head [0 [[[[3]]]] [[[2]]] [[1]] [0]] ⧗ (Digit a) → a

:test (digit-head (one i)) (i)
:test (digit-head (two i k)) (i)
:test (digit-head (three i k k)) (i)
:test (digit-head (four i k k k)) (i)

# returns trailing elemente of digit
digit-tail [0 [three] [two] [one] Ω] ⧗ (Digit a) → (Digit a)

:test (digit-tail (two i k)) (one k)
:test (digit-tail (three i k k)) (two k k)
:test (digit-tail (four i k k k)) (three k k k)

foldr-digit [[[0 case-four case-three case-two case-one]]] ⧗ (a → b → b) → b → (Digit a) → b
	case-four [[[[6 3 (6 2 (6 1 (6 0 5)))]]]]
	case-three [[[5 2 (5 1 (5 0 4))]]]
	case-two [[4 1 (4 0 3)]]
	case-one [3 0 2]

:test (foldr-digit L.cons L.empty (one 'a')) ("a")
:test (foldr-digit L.cons L.empty (two 'a' 'b')) ("ab")
:test (foldr-digit L.cons L.empty (three 'a' 'b' 'c')) ("abc")
:test (foldr-digit L.cons L.empty (four 'a' 'b' 'c' 'd')) ("abcd")

foldl-digit [[[0 case-four case-three case-two case-one]]] ⧗ (a → b → b) → b → (Digit a) → b
	case-four [[[[6 (6 (6 (6 5 3) 2) 1) 0]]]]
	case-three [[[5 (5 (5 4 2) 1) 0]]]
	case-two [[4 (4 3 1) 0]]
	case-one [3 2 0]

:test (foldl-digit \L.cons L.empty (one 'a')) ("a")
:test (foldl-digit \L.cons L.empty (two 'a' 'b')) ("ba")
:test (foldl-digit \L.cons L.empty (three 'a' 'b' 'c')) ("cba")
:test (foldl-digit \L.cons L.empty (four 'a' 'b' 'c' 'd')) ("dcba")

# adds element to digit (and updates its tag)
cons [[0 case-four case-three case-two case-one]] ⧗ a → (Digit a) → (Digit a)
	case-four Ω
	case-three [[[[[[[3 8 6 5 4]]]]]]]
	case-two [[[[[[2 7 5 4]]]]]]
	case-one [[[[[1 6 4]]]]]

…:… cons

:test ('a' : (one 'b')) (two 'a' 'b')
:test ('a' : (two 'b' 'c')) (three 'a' 'b' 'c')
:test ('a' : (three 'b' 'c' 'd')) (four 'a' 'b' 'c' 'd')

# adds element to digit (and updates its tag)
snoc [[1 case-four case-three case-two case-one]] ⧗ (Digit a) → a → (Digit a)
	case-four Ω
	case-three [[[[[[[3 6 5 4 7]]]]]]]
	case-two [[[[[[2 5 4 6]]]]]]
	case-one [[[[[1 4 5]]]]]

…;… snoc

:test ((one 'a') ; 'b') (two 'a' 'b')
:test ((two 'a' 'b') ; 'c') (three 'a' 'b' 'c')
:test ((three 'a' 'b' 'c') ; 'd') (four 'a' 'b' 'c' 'd')

# removes redundant abstractions
↓₃‣ [0 i i i]

# === Tree ===
# Scott-style tagged union, 3 tags

# (FingerTree a) = Empty
#                | (Single a)
#                | (Deep (Digit a) (FingerTree (Node a)) (Digit a))

# empty tree (tag 0)
empty [[[0]]] ⧗ (FingerTree a)

# single tree (tag 1)
single [[[[1 3]]]] ⧗ a → (FingerTree a)

# deep tree (tag 2)
deep [[[[[[2 5 4 3]]]]]] ⧗ (Digit a) → (FingerTree (Node a)) → (Digit a) → (FingerTree a)

foldr-tree z [[[[0 case-deep case-single case-empty]]]] ⧗ (a → b → b) → b → (FingerTree a) → b
	case-deep [[[foldr-digit 5 (6 \(foldr-node 5) (foldr-digit 5 4 0) 1) 2]]]
	case-single [3 0 2]
	case-empty 1

foldl-tree z [[[[0 case-deep case-single case-empty]]]] ⧗ (b → a → b) → b → (FingerTree a) → b
	case-deep [[[foldl-digit 5 (6 (foldl-node 5) (foldl-digit 5 4 2) 1) 0]]]
	case-single [3 2 0]
	case-empty 1

# adds element to the left side of a finger tree
insert-left z [[[0 case-deep case-single case-empty]]] ⧗ a → (FingerTree a) → (FingerTree a)
	case-deep [[[four? 2 overflow append]]]
		overflow deep (two 4 ↓₃(2 [[[[3]]]])) (5 ↓₃(2 [node3]) 1) 0
		append deep (4 : 2) 1 0
	case-single [deep {}2 empty {}0]
	case-empty single 1

…◁… insert-left

:test ('a' ◁ empty) (single 'a')
:test ('a' ◁ (single 'b')) (deep (one 'a') empty (one 'b'))
:test ('a' ◁ (deep (three 'b' 'c' 'd') empty (one 'a'))) (deep (four 'a' 'b' 'c' 'd') empty (one 'a'))
:test ('a' ◁ (deep (four 'b' 'c' 'd' 'e') empty (one 'a'))) (deep (two 'a' 'b') (single (node3 'c' 'd' 'e')) (one 'a'))

……◁′… [\(0 insert-left)] ⧗ (Foldr s) → (s a) → (FingerTree a) → (FingerTree a)

:test (L.foldr "abcdefg" ◁′ empty) (deep (three 'a' 'b' 'c') (single (node3 'd' 'e' 'f')) (one 'g'))

# adds element to the right side of a finger tree
insert-right z [[[1 case-deep case-single case-empty]]] ⧗ (FingerTree a) → a → (FingerTree a)
	case-deep [[[four? 0 overflow append]]]
		overflow deep 2 (5 1 ↓₃(0 [[[[node3 3 2 1]]]])) (two ↓₃(0 [[[[0]]]]) 3)
		append deep 2 1 (0 ; 3)
	case-single [deep {}0 empty {}1]
	case-empty single 0

…▷… insert-right

:test (empty ▷ 'a') (single 'a')
:test ((single 'a') ▷ 'b') (deep (one 'a') empty (one 'b'))
:test ((deep (one 'a') empty (three 'a' 'b' 'c')) ▷ 'd') (deep (one 'a') empty (four 'a' 'b' 'c' 'd'))
:test ((deep (one 'a') empty (four 'e' 'd' 'c' 'b')) ▷ 'a') (deep (one 'a') (single (node3 'e' 'd' 'c')) (two 'b' 'a'))

……▷′… &insert-right ⧗ (Foldl s) → (FingerTree a) → (s a) → (FingerTree a)

:test (L.foldl empty ▷′ "abcdefg") (deep (one 'a') (single (node3 'b' 'c' 'd')) (three 'e' 'f' 'g'))

# === Conversions ===

# converts a list to a finger tree
list→tree [L.foldr 0 ◁′ empty] ⧗ (List a) → (FingerTree a)

:test (list→tree "a") (single 'a')
:test (list→tree "abcdefu") (deep (three 'a' 'b' 'c') (single (node3 'd' 'e' 'f')) (one 'u'))
:test (foldl-tree \L.cons L.empty (list→tree "abcdefghijklmnopqrstuvwxyz")) ("zyxwvutsrqponmlkjihgfedcba")

# converts a digit to a finger tree
digit→tree [foldr-digit 0 ◁′ empty] ⧗ (Digit a) → (FingerTree a)

# converts a digit to a list
digit→list foldr-digit L.cons L.empty ⧗ (Digit a) → (List a)

# converts a node to a digit
node→digit [0 three two] ⧗ (Node a) → (Digit a)

:test (node→digit (node2 'a' 'b')) (two 'a' 'b')
:test (node→digit (node3 'a' 'b' 'c')) (three 'a' 'b' 'c')

# converts a finger tree to a list
tree→list foldr-tree L.cons L.empty ⧗ (FingerTree a) → (List a)

:test (tree→list (list→tree "a")) ("a")
:test (tree→list (list→tree "in ulm, um ulm und um ulm herum")) ("in ulm, um ulm und um ulm herum")

# === View ===

# (ViewL a) = Empty | Pair a (FingerTree a)

# constructs a view of the tree
# basically shifts the leftmost element into a pair
view-left z [[0 case-deep case-single case-empty]] ⧗ (FingerTree a) → (ViewL a)
	case-deep [[[L.cons (digit-head 2) (2 [deep ∘∘∘ three] [deep ∘∘ two] [deep ∘ one] [deep-left] 1 0)]]]
		deep-left [[7 1 [[[case-cons]]] case-nil]] ⧗ (FingerTree (Node a)) → (Digit a) → (FingerTree a)
			case-cons deep (node→digit 2) 1 3
			case-nil digit→tree 0
	case-single [L.cons 0 empty]
	case-empty L.empty

:test (view-left empty) (L.empty)
:test (view-left (single 'a')) (L.cons 'a' empty)
:test (view-left (deep (two 'a' 'b') empty (one 'c'))) (L.cons 'a' (deep (one 'b') empty (one 'c')))
:test (view-left (deep (one 'a') empty (two 'b' 'c'))) (L.cons 'a' (deep (one 'b') empty (one 'c')))
:test (view-left (deep (one 'a') (single (node3 'b' 'c' 'd')) (one 'e'))) (L.cons 'a' (deep (three 'b' 'c' 'd') empty (one 'e')))

# returns true if finger tree is empty
empty? [view-left 0 [[[ki]]] k] ⧗ (FingerTree a) → Boolean

:test (empty? empty) (k)
:test (empty? (single 'a')) (ki)
:test (empty? (deep (one 'a') empty (one 'b'))) (ki)

# returns left head of finger tree
head-left L.head ∘ view-left ⧗ (FingerTree a) → a

:test (head-left (list→tree "abcdefg")) ('a')

# returns left tail of finger tree
tail-left L.tail ∘ view-left ⧗ (FingerTree a) → (FingerTree a)

# TODO: implement viewR (mirror image)

# === Concatenation ===

# WARNING: this will only work for lengths with factor 2 or 3
# case-+ is also not really relevant I think
list→nodes [z [[[rec]]] 0 L.∀0] ⧗ (List a) → (List (Node a))
	rec N.eq? 0 (+2) case-2 (N.eq? 0 (+3) case-3 (N.eq? 0 (+4) case-4 case-+))
		case-2 1 [[L.{}(node2 1 L.^0)]]
		case-3 1 [[0 [[L.{}(node3 3 1 L.^0)]]]]
		case-4 1 [[0 [[0 [[L.cons (node2 5 3) L.{}(node2 1 L.^0)]]]]]]
		case-+ 1 [[0 [[0 [[L.cons (node3 5 3 1) (8 0 (N.sub 6 (+3)))]]]]]]

:test (list→nodes "ab") (L.{}(node2 'a' 'b'))
:test (list→nodes "abc") (L.{}(node3 'a' 'b' 'c'))
:test (list→nodes "abcd") (L.cons (node2 'a' 'b') L.{}(node2 'c' 'd'))
:test (list→nodes "abcde") (L.cons (node3 'a' 'b' 'c') L.{}(node2 'd' 'e'))

append3 z [[[[2 case-deep case-single case-empty]]]] ⧗ (FingerTree a) → (List a) → (FingerTree a) → (FingerTree a)
	case-deep [[[3 deep-deep deep-single deep-empty]]]
		deep-deep [[[deep 5 (9 4 new-list 1) 0]]]
			new-list list→nodes (L.append (L.append (digit→list 3) 7) (digit→list 2))
		deep-single [(L.foldl 6 ▷′ 5) ▷ 0]
		deep-empty 5
	case-single [1 single-deep single-single single-empty]
		single-deep [[[3 ◁ (L.foldr 5 ◁′ 4)]]]
		single-single [1 ◁ (L.foldr 3 ◁′ 2)]
		single-empty 3
	case-empty 0

append [[append3 1 L.empty 0]] ⧗ (FingerTree a) → (FingerTree a) → (FingerTree a)

…++… append

:test (tree→list ((list→tree "a") ++ (list→tree "b"))) ("ab")
:test (tree→list ((list→tree "abcdefg") ++ (list→tree "hijklmnop"))) ("abcdefghijklmnop")
:test (tree→list ((list→tree "abcdefghijklmnopqrstuvwxyz1234") ++ (list→tree "abcdefghijklmopqrstuvwxyz"))) ("abcdefghijklmnopqrstuvwxyz1234abcdefghijklmopqrstuvwxyz")

# TODO: annotations, measurement, splitting, random-access
#   - annotations will require some modifications (more abstractions)
# TODO: new modules: sequence, pqueue