Files
newt/aoc2024/SortedMap.newt
2024-12-06 09:37:20 -08:00

76 lines
3.0 KiB
Agda
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
module SortedMap
import Prelude
data T23 : Nat -> U -> U -> U where
Leaf : k v. k -> v -> T23 Z k v
Node2 : h k v. T23 h k v -> k -> T23 h k v -> T23 (S h) k v
Node3 : h k v. T23 h k v -> k -> T23 h k v -> k -> T23 h k v -> T23 (S h) k v
lookupT23 : h k v. {{Ord k}} {{Eq k}} -> k -> T23 h k v -> Maybe (k × v)
lookupT23 key (Leaf k v)= if k == key then Just (k,v) else Nothing
lookupT23 key (Node2 t1 k1 t2) =
if key <= k1 then lookupT23 key t1 else lookupT23 key t2
lookupT23 key (Node3 t1 k1 t2 k2 t3) =
if key <= k1 then lookupT23 key t1
else if key <= k2 then lookupT23 key t2
else lookupT23 key t3
insertT23 : h k v. {{Ord k}} {{Eq k}} -> k -> v -> T23 h k v -> Either (T23 h k v) (T23 h k v × k × T23 h k v)
insertT23 key value (Leaf k v) =
if key == k then Left (Leaf key value)
else if key <= k then Right (Leaf key value, key, Leaf k v)
else Right (Leaf k v, k, Leaf key value)
insertT23 key value (Node2 t1 k1 t2) =
if key <= k1 then
case insertT23 key value t1 of
Left t1' => Left (Node2 t1' k1 t2)
Right (a,b,c) => Left (Node3 a b c k1 t2)
else case insertT23 key value t2 of
Left t2' => Left (Node2 t1 k1 t2')
Right (a,b,c) => Left (Node3 t1 k1 a b c)
insertT23 key value (Node3 t1 k1 t2 k2 t3) =
if key <= k1 then
case insertT23 key value t1 of
Left t1' => Left (Node3 t1' k1 t2 k2 t3)
Right (a,b,c) => Right (Node2 a b c, k1, Node2 t2 k2 t3)
else if key <= k2 then
case insertT23 key value t2 of
Left t2' => Left (Node3 t1 k1 t2' k2 t3)
Right (a,b,c) => Right (Node2 t1 k1 a, b, Node2 c k2 t3)
else
case insertT23 key value t3 of
Left t3' => Left (Node3 t1 k1 t2 k2 t3')
Right (a,b,c) => Right (Node2 t1 k1 t2, k2, Node2 a b c)
data SortedMap : U -> U -> U where
EmptyMap : k v. SortedMap k v
MapOf : k v h. T23 h k v -> SortedMap k v
lookupMap : k v. {{Ord k}} {{Eq k}} -> k -> SortedMap k v -> Maybe (k × v)
lookupMap k EmptyMap = Nothing
lookupMap k (MapOf map) = lookupT23 k map
updateMap : k v. {{Ord k}} {{Eq k}} -> k -> v -> SortedMap k v -> SortedMap k v
updateMap k v EmptyMap = MapOf $ Leaf k v
updateMap k v (MapOf map) = case insertT23 k v map of
Left map' => MapOf map'
Right (a, b, c) => MapOf (Node2 a b c)
-- FIXME this doesn't work in a `where` because the erased args are un-erased
toList' : k v h. T23 h k v List (k × v) List (k × v)
toList' (Leaf k v) acc = (k, v) :: acc
toList' (Node2 t1 k1 t2) acc = toList' t2 (toList' t1 acc)
toList' (Node3 t1 k1 t2 k2 t3) acc = toList' t3 $ toList' t2 $ toList' t1 acc
toList : k v. SortedMap k v List (k × v)
toList {k} {v} (MapOf smap) = reverse $ toList' smap Nil
-- FIXME erasure checking false positive - maybe because I'm not handling the top level args yet
-- where
-- foo : ∀ k v h. T23 h k v → List (k × v) → List (k × v)
-- foo (Leaf k v) acc = (k, v) :: acc
-- foo (Node2 t1 k1 t2) acc = foo t2 (foo t1 acc)
-- foo (Node3 t1 k1 t2 k2 t3) acc = foo t3 $ foo t2 $ foo t1 acc
toList _ = Nil