Files
newt/aoc2024/SortedMap.newt

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 Letrec doesn't have the type
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 (need type on foo in env when processing Letrec)
-- 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