Map2.agda

module Map2 where

--------------------------------------------------
-- * Prelude
--------------------------------------------------
-- Taken from: https://github.com/pigworker/Ohrid-Agda
open import Ohrid-Prelude
open import Ohrid-Nat

-- Mini HoTT (Except that we have K)
ap : ∀ {A B : Set} {x y : A} (f : A → B) (p : x == y) → (f x == f y)
ap f refl = refl

transport : {A : Set} (B : A -> Set) {a b : A} (p : a == b) -> B a -> B b
transport B refl = id

data Σ (A : Set) (B : A -> Set) : Set where
  sig : (a : A) -> (b : B a) -> Σ A B

data Maybe (A : Set) : Set where
  Nothing : Maybe A
  Just    : A -> Maybe A

data Bool : Set where
  true  : Bool
  false : Bool

isSome : {A : Set} -> Maybe A -> Bool
isSome Nothing = false
isSome (Just _) = true

negb : Bool -> Bool
negb true = false
negb false = true

bleq : Bool -> Bool -> Set
bleq true false = Zero
bleq _    _     = One

bleq-uniq : {b1 b2 : Bool} (p q : bleq b1 b2) -> p == q
bleq-uniq {true} {true} <> <> = refl
bleq-uniq {true} {false} () q
bleq-uniq {false} {b2} <> <> = refl

bleq-t : {b : Bool} -> bleq b true
bleq-t {true} = <>
bleq-t {false} = <>

--------------------------------------------------
-- Pre maps
--------------------------------------------------

data PreMap (A : Set) : Bool {- is the map nil? -} → Set where
  mnil  : PreMap A true
  mcons : (x : Maybe A) {b : Bool} (w : bleq b (isSome x)) (m : PreMap A b) -> PreMap A false

prelookup : {A : Set} {b : Bool} -> PreMap A b -> Nat -> Maybe A
prelookup mnil i = Nothing
prelookup (mcons x w m) zero = x
prelookup (mcons x w m) (suc i) = prelookup m i

premap-wf-nil : forall {A : Set} {b : Bool} (m : PreMap A b) ->
  (forall (i : Nat) -> prelookup m i == Nothing) -> Σ (b == true) (λ p -> transport (PreMap A) p m == mnil)
premap-wf-nil mnil p = sig refl refl
premap-wf-nil (mcons x w m) p with p 0
premap-wf-nil (mcons .Nothing w m) p | refl with premap-wf-nil m (λ i -> p (suc i))
premap-wf-nil (mcons .Nothing () m) p | refl | sig refl _

premap-prelookup-eq : {A : Set} {b1 b2 : Bool} -> (m1 : PreMap A b1) (m2 : PreMap A b2) -> 
  (forall (i : Nat) -> prelookup m1 i == prelookup m2 i) -> Σ (b1 == b2) (λ p -> transport (PreMap A) p m1 == m2)
premap-prelookup-eq m1 mnil p = premap-wf-nil m1 p
premap-prelookup-eq mnil (mcons x w m2) p with premap-wf-nil (mcons x w m2) (λ i -> sym (p i))
premap-prelookup-eq mnil (mcons x w m2) p | sig () _
premap-prelookup-eq (mcons x w m1) (mcons y v m2) p with p 0
premap-prelookup-eq (mcons x w m1) (mcons .x v m2) p | refl
   with premap-prelookup-eq m1 m2 (λ i -> p (suc i))
premap-prelookup-eq (mcons x w m1) (mcons .x v m2) p | refl | sig refl q
  rewrite bleq-uniq w v | q = sig refl refl   

presingleton : {A : Set} (i : Nat) (v : A) -> PreMap A false
presingleton zero v = mcons (Just v) <> mnil 
presingleton (suc i) v = mcons Nothing <> (presingleton i v)

preinsert : {b : Bool} {A : Set} (i : Nat) (v : A) -> PreMap A b -> PreMap A false
preinsert i v mnil = presingleton i v
preinsert zero v (mcons x w m) = mcons (Just v) bleq-t m
preinsert (suc i) v (mcons x w m) = mcons x <> (preinsert i v m)


------------------------------------------------
-- Actual maps
------------------------------------------------

data Map (A : Set) : Set where
  M : {b : Bool} (m : PreMap A b) -> Map A

mempty : {A : Set} -> Map A
mempty = M mnil

_!!_ : {A : Set} -> Map A -> Nat -> Maybe A
M m !! i = prelookup m i

mlookup-canon : {A : Set} (m1 m2 : Map A) ->
  (forall (i : Nat) -> m1 !! i == m2 !! i) -> m1 == m2
mlookup-canon (M m1) (M m2) p with premap-prelookup-eq m1 m2 p
mlookup-canon (M m1) (M m2) p | sig refl q = ap M q

[_:=_] : {A : Set} (i : Nat) (v : A) -> Map A
[ i := v ] = M (presingleton i v)

minsert : {A : Set} (i : Nat) (v : A) -> Map A -> Map A
minsert i v (M m) = M (preinsert i v m)

--------------------------------------------------
-- Properties
--------------------------------------------------

singleton-lookup : {A : Set} (i : Nat) (v : A) -> [ i := v ] !! i == Just v
singleton-lookup zero v = refl
singleton-lookup (suc i) v = singleton-lookup i v

insert-lookup : {A : Set} (i : Nat) (v : A) (m : Map A) -> (minsert i v m) !! i == Just v
insert-lookup i v (M mnil) = singleton-lookup i v
insert-lookup zero v (M (mcons x w m)) = refl
insert-lookup (suc i) v (M (mcons x w m)) = insert-lookup i v (M m)

mempty-lookup : {A : Set} (i : Nat) -> mempty {A} !! i == Nothing
mempty-lookup i = refl

map-fmap-raw : {A B : Set} {b : Bool} -> (A -> B) -> PreMap A b -> PreMap B b
map-fmap-raw f mnil = mnil
map-fmap-raw f (mcons Nothing w m) = mcons Nothing w (map-fmap-raw f m)
map-fmap-raw f (mcons (Just x) w m) = mcons (Just (f x)) w (map-fmap-raw f m)

map-fmap : {A B : Set} -> (A -> B) -> Map A -> Map B
map-fmap f (M m) = M (map-fmap-raw f m)