1 files changed, 138 insertions, 0 deletions
diff --git a/src/CHAD/AST/Weaken.hs b/src/CHAD/AST/Weaken.hs
new file mode 100644
index 0000000..ac0d152
--- /dev/null
+++ b/src/CHAD/AST/Weaken.hs
@@ -0,0 +1,138 @@
+{-# LANGUAGE DataKinds #-}
+{-# LANGUAGE EmptyCase #-}
+{-# LANGUAGE GADTs #-}
+{-# LANGUAGE PolyKinds #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE StandaloneDeriving #-}
+{-# LANGUAGE StandaloneKindSignatures #-}
+{-# LANGUAGE TypeAbstractions #-}
+{-# LANGUAGE TypeApplications #-}
+{-# LANGUAGE TypeOperators #-}
+
+{-# LANGUAGE PartialTypeSignatures #-}
+{-# OPTIONS -Wno-partial-type-signatures #-}
+
+-- The reason why this is a separate module with "little" in it:
+{-# LANGUAGE AllowAmbiguousTypes #-}
+
+module CHAD.AST.Weaken (module CHAD.AST.Weaken, Append) where
+
+import Data.Bifunctor (first)
+import Data.Functor.Const
+import Data.GADT.Compare
+import Data.Kind (Type)
+
+import CHAD.Data
+import CHAD.Lemmas
+
+
+type Idx :: [k] -> k -> Type
+data Idx env t where
+  IZ :: Idx (t : env) t
+  IS :: Idx env t -> Idx (a : env) t
+deriving instance Show (Idx env t)
+
+instance GEq (Idx env) where
+  geq IZ IZ = Just Refl
+  geq (IS i) (IS j) | Just Refl <- geq i j = Just Refl
+  geq _ _ = Nothing
+
+splitIdx :: forall env2 env1 t f. SList f env1 -> Idx (Append env1 env2) t -> Either (Idx env1 t) (Idx env2 t)
+splitIdx SNil i = Right i
+splitIdx (SCons _ _) IZ = Left IZ
+splitIdx (SCons _ l) (IS i) = first IS (splitIdx l i)
+
+slistIdx :: SList f list -> Idx list t -> f t
+slistIdx (SCons x _) IZ = x
+slistIdx (SCons _ list) (IS i) = slistIdx list i
+slistIdx SNil i = case i of {}
+
+idx2int :: Idx env t -> Int
+idx2int IZ = 0
+idx2int (IS n) = 1 + idx2int n
+
+data env :> env' where
+  WId :: env :> env
+  WSink :: forall t env. env :> (t : env)
+  WCopy :: forall t env env'. env :> env' -> (t : env) :> (t : env')
+  WPop :: (t : env) :> env' -> env :> env'
+  WThen :: env1 :> env2 -> env2 :> env3 -> env1 :> env3
+  WClosed :: '[] :> env
+  WIdx :: Idx env t -> (t : env) :> env
+  WPick :: forall t pre env env'. SList (Const ()) pre -> env :> env'
+                               -> Append pre (t : env) :> t : Append pre env'
+  WSwap :: forall env as bs. SList (Const ()) as -> SList (Const ()) bs
+        -> Append as (Append bs env) :> Append bs (Append as env)
+  WStack :: forall env1 env2 as bs. SList (Const ()) as -> SList (Const ()) bs
+         -> as :> bs -> env1 :> env2
+         -> Append as env1 :> Append bs env2
+deriving instance Show (env :> env')
+infix 4 :>
+
+infixr 2 @>
+(@>) :: env :> env' -> Idx env t -> Idx env' t
+WId @> i = i
+WSink @> i = IS i
+WCopy _ @> IZ = IZ
+WCopy w @> IS i = IS (w @> i)
+WPop w @> i = w @> IS i
+WThen w1 w2 @> i = w2 @> w1 @> i
+WClosed @> i = case i of {}
+WIdx j @> IZ = j
+WIdx _ @> IS i = i
+WPick SNil w @> i = WCopy w @> i
+WPick (_ `SCons` _) _ @> IZ = IS IZ
+WPick @t (_ `SCons` pre) w @> IS i = WCopy WSink .> WPick @t pre w @> i
+WSwap @env (as :: SList _ as) (bs :: SList _ bs) @> i =
+  case splitIdx @(Append bs env) as i of
+    Left i' -> indexSinks bs (indexRaiseAbove @env as i')
+    Right i' -> case splitIdx @env bs i' of
+                  Left j -> indexRaiseAbove @(Append as env) bs j
+                  Right j -> indexSinks bs (indexSinks as j)
+WStack @env1 @env2 as bs wlo whi @> i =
+  case splitIdx @env1 as i of
+    Left i' -> indexRaiseAbove @env2 bs (wlo @> i')
+    Right i' -> indexSinks bs (whi @> i')
+
+indexSinks :: SList f as -> Idx bs t -> Idx (Append as bs) t
+indexSinks SNil j = j
+indexSinks (_ `SCons` bs') j = IS (indexSinks bs' j)
+
+indexRaiseAbove :: forall env as t f. SList f as -> Idx as t -> Idx (Append as env) t
+indexRaiseAbove = flip go
+  where
+    go :: forall as'. Idx as' t -> SList f as' -> Idx (Append as' env) t
+    go IZ (_ `SCons` _) = IZ
+    go (IS i) (_ `SCons` as) = IS (go i as)
+
+infixr 3 .>
+(.>) :: env2 :> env3 -> env1 :> env2 -> env1 :> env3
+(.>) = flip WThen
+
+class KnownListSpine list where knownListSpine :: SList (Const ()) list
+instance KnownListSpine '[] where knownListSpine = SNil
+instance KnownListSpine list => KnownListSpine (t : list) where knownListSpine = SCons (Const ()) knownListSpine
+
+wSinks' :: forall list env. KnownListSpine list => env :> Append list env
+wSinks' = wSinks (knownListSpine :: SList (Const ()) list)
+
+wSinks :: forall env bs f. SList f bs -> env :> Append bs env
+wSinks SNil = WId
+wSinks (SCons _ spine) = WSink .> wSinks spine
+
+wSinksAnd :: forall env env' bs f. SList f bs -> env :> env' -> env :> Append bs env'
+wSinksAnd SNil w = w
+wSinksAnd (SCons _ spine) w = WSink .> wSinksAnd spine w
+
+wCopies :: SList f bs -> env1 :> env2 -> Append bs env1 :> Append bs env2
+wCopies bs w =
+  let bs' = slistMap (\_ -> Const ()) bs
+  in WStack bs' bs' WId w
+
+wRaiseAbove :: SList f env1 -> proxy env -> env1 :> Append env1 env
+wRaiseAbove SNil _ = WClosed
+wRaiseAbove (SCons _ s) env = WCopy (wRaiseAbove s env)
+
+wPops :: SList f bs -> Append bs env1 :> env2 -> env1 :> env2
+wPops SNil w = w
+wPops (_ `SCons` bs) w = wPops bs (WPop w)