Towards neural

2 files changed, 83 insertions, 34 deletions

diff --git a/src/AST/Weaken.hs b/src/AST/Weaken.hs
index 78276ca..0a1e4ce 100644
--- a/src/AST/Weaken.hs
+++ b/src/AST/Weaken.hs
@@ -48,6 +48,9 @@ data env :> env' where
                                -> 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 as bs env1 env2. SList (Const ()) as -> SList (Const ()) bs
+         -> as :> bs -> env1 :> env2
+         -> Append as env1 :> Append bs env2
 deriving instance Show (env :> env')
 infix 4 :>
 
@@ -67,18 +70,25 @@ 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' -> skipOver bs (stack @env i' as)
+    Left i' -> indexSinks bs (indexRaiseAbove @env as i')
     Right i' -> case splitIdx @env bs i' of
-                  Left j -> stack @(Append as env) j bs
-                  Right j -> skipOver bs (skipOver as j)
+                  Left j -> indexRaiseAbove @(Append as env) bs j
+                  Right j -> indexSinks bs (indexSinks as j)
+WStack @as @bs @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
-    skipOver :: SList (Const ()) as' -> Idx bs' t -> Idx (Append as' bs') t
-    skipOver SNil j = j
-    skipOver (_ `SCons` bs') j = IS (skipOver bs' j)
-
-    stack :: forall env' as' t. Idx as' t -> SList (Const ()) as' -> Idx (Append as' env') t
-    stack IZ (_ `SCons` _) = IZ
-    stack (IS j) (_ `SCons` as') = IS (stack @env' j as')
+    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
@@ -100,8 +110,9 @@ 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 SNil w = w
-wCopies (SCons _ spine) w = WCopy (wCopies spine w)
+wCopies bs w =
+  let bs' = slistMap (\_ -> Const ()) bs
+  in WStack bs' bs' WId w
 
 wRaiseAbove :: SList f env1 -> SList g env -> env1 :> Append env1 env
 wRaiseAbove SNil env = WClosed (slistMap (\_ -> Const ()) env)
diff --git a/src/AST/Weaken/Auto.hs b/src/AST/Weaken/Auto.hs
index 444c540..8555516 100644
--- a/src/AST/Weaken/Auto.hs
+++ b/src/AST/Weaken/Auto.hs
@@ -35,15 +35,29 @@ import Lemmas
 type family Lookup name list where
   Lookup name ('(name, x) : _) = x
   Lookup name (_ : list) = Lookup name list
+  Lookup name '[] = TypeError (Text "The name '" :<>: Text name :<>: Text "' does not appear in the list.")
+
+
+-- | The @withPre@ type parameter indicates whether there can be 'LPreW'
+-- occurrences within this layout.
+data Layout (withPre :: Bool) (segments :: [(Symbol, [t])]) (env :: [t]) where
+  LSeg :: forall name segments withPre. SSymbol name -> Layout withPre segments (Lookup name segments)
+  -- | Pre-weaken with a weakening
+  LPreW :: forall name1 name2 segments.
+           SegmentName name1 -> SegmentName name2
+        -> Lookup name1 segments :> Lookup name2 segments
+        -> Layout True segments (Lookup name1 segments)
+  (:++:) :: Layout withPre segments env1 -> Layout withPre segments env2 -> Layout withPre segments (Append env1 env2)
+infixr :++:
 
+instance (KnownSymbol name, seg ~ Lookup name segments) => IsLabel name (Layout withPre segments seg) where
+  fromLabel = LSeg (symbolSing @name)
 
-data Layout (segments :: [(Symbol, [t])]) (env :: [t]) where
-  LSeg :: forall name segments. KnownSymbol name => Layout segments (Lookup name segments)
-  (:++:) :: Layout segments env1 -> Layout segments env2 -> Layout segments (Append env1 env2)
-infixr :++:
+newtype SegmentName name = SegmentName (SSymbol name)
+  deriving (Show)
 
-instance (KnownSymbol name, seg ~ Lookup name segments) => IsLabel name (Layout segments seg) where
-  fromLabel = LSeg @name @segments
+instance (KnownSymbol name, name ~ name') => IsLabel name (SegmentName name') where
+  fromLabel = SegmentName symbolSing
 
 
 data SSegments (segments :: [(Symbol, [t])]) where
@@ -86,29 +100,51 @@ segmentLookup = \segs name -> case go segs name of
           case unsafeCoerce Refl :: (Lookup name ('(n, ts) : rest) :~: Lookup name rest) of
             Refl -> go sseg name
 
-data LinLayout (segments :: [(Symbol, [t])]) (env :: [t]) where
-  LinEnd :: LinLayout segments '[]
-  LinApp :: SSymbol name -> LinLayout segments env -> LinLayout segments (Append (Lookup name segments) env)
+data LinLayout (withPre :: Bool) (segments :: [(Symbol, [t])]) (env :: [t]) where
+  LinEnd :: LinLayout withPre segments '[]
+  LinApp :: SSymbol name -> LinLayout withPre segments env
+         -> LinLayout withPre segments (Append (Lookup name segments) env)
+  LinAppPreW :: SSymbol name1 -> SSymbol name2
+             -> Lookup name1 segments :> Lookup name2 segments
+             -> LinLayout True segments env
+             -> LinLayout True segments (Append (Lookup name1 segments) env)
 
-linLayoutAppend :: LinLayout segments env1 -> LinLayout segments env2 -> LinLayout segments (Append env1 env2)
+linLayoutAppend :: LinLayout withPre segments env1 -> LinLayout withPre segments env2 -> LinLayout withPre segments (Append env1 env2)
 linLayoutAppend LinEnd lin = lin
-linLayoutAppend (LinApp (name :: SSymbol name) (lin1 :: LinLayout segments env1')) (lin2 :: LinLayout _ env2)
+linLayoutAppend (LinApp (name :: SSymbol name) (lin1 :: LinLayout _ segments env1')) (lin2 :: LinLayout _ _ env2)
   | Refl <- lemAppendAssoc @(Lookup name segments) @env1' @env2
   = LinApp name (linLayoutAppend lin1 lin2)
+linLayoutAppend (LinAppPreW (name1 :: SSymbol name1) name2 w (lin1 :: LinLayout _ segments env1')) (lin2 :: LinLayout _ _ env2)
+  | Refl <- lemAppendAssoc @(Lookup name1 segments) @env1' @env2
+  = LinAppPreW name1 name2 w (linLayoutAppend lin1 lin2)
 
-linLayoutEnv :: SSegments segments -> LinLayout segments env -> SList (Const ()) env
+linLayoutEnv :: SSegments segments -> LinLayout withPre segments env -> SList (Const ()) env
 linLayoutEnv _ LinEnd = SNil
 linLayoutEnv segs (LinApp name lin) = sappend (segmentLookup segs name) (linLayoutEnv segs lin)
+linLayoutEnv segs (LinAppPreW name1 _ _ lin) = sappend (segmentLookup segs name1) (linLayoutEnv segs lin)
 
-lineariseLayout :: Layout segments env -> LinLayout segments env
-lineariseLayout (LSeg @name :: Layout _ seg)
+lineariseLayout :: Layout withPre segments env -> LinLayout withPre segments env
+lineariseLayout (LSeg name :: Layout _ _ seg)
   | Refl <- lemAppendNil @seg
-  = LinApp (symbolSing @name) LinEnd
+  = LinApp name LinEnd
 lineariseLayout (ly1 :++: ly2) = lineariseLayout ly1 `linLayoutAppend` lineariseLayout ly2
-
-pullDown :: SSegments segments -> SSymbol name -> LinLayout segments env
+lineariseLayout (LPreW (SegmentName name1) (SegmentName name2) w :: Layout _ _ seg)
+  | Refl <- lemAppendNil @seg
+  = LinAppPreW name1 name2 w LinEnd
+
+preWeaken :: SSegments segments -> LinLayout True segments env
+          -> (forall env'. env :> env' -> LinLayout False segments env' -> r) -> r
+preWeaken _ LinEnd k = k WId LinEnd
+preWeaken segs (LinApp name lin) k =
+  preWeaken segs lin $ \w lin' ->
+    k (wCopies (segmentLookup segs name) w) (LinApp name lin')
+preWeaken segs (LinAppPreW name1 name2 weak lin) k =
+  preWeaken segs lin $ \w lin' ->
+    k (WStack (segmentLookup segs name1) (segmentLookup segs name2) weak w) (LinApp name2 lin')
+
+pullDown :: SSegments segments -> SSymbol name -> LinLayout False segments env
          -> r  -- Name was not found in source
-         -> (forall env'. LinLayout segments env' -> env :> Append (Lookup name segments) env' -> r)
+         -> (forall env'. LinLayout False segments env' -> env :> Append (Lookup name segments) env' -> r)
          -> r
 pullDown segs name@SSymbol linlayout kNotFound k =
   case linlayout of
@@ -116,13 +152,13 @@ pullDown segs name@SSymbol linlayout kNotFound k =
     LinApp n'@SSymbol lin
       | Just Refl <- sameSymbol name n' -> k lin WId
       | otherwise ->
-          pullDown segs name lin kNotFound $ \(lin' :: LinLayout _ env') w ->
+          pullDown segs name lin kNotFound $ \(lin' :: LinLayout _ _ env') w ->
             k (LinApp n' lin') (WSwap @env' (segmentLookup segs n') (segmentLookup segs name)
                                   .> wCopies (segmentLookup segs n') w)
 
 sortLinLayouts :: forall segments env1 env2.
                   SSegments segments
-               -> LinLayout segments env1 -> LinLayout segments env2 -> env1 :> env2
+               -> LinLayout False segments env1 -> LinLayout False segments env2 -> env1 :> env2
 sortLinLayouts _ LinEnd LinEnd = WId
 sortLinLayouts segs lin1@(LinApp name1@SSymbol tail1) (LinApp name2@SSymbol tail2)
   | Just Refl <- sameSymbol name1 name2 = wCopies (segmentLookup segs name1) (sortLinLayouts segs tail1 tail2)
@@ -139,5 +175,7 @@ sortLinLayouts segs LinEnd lin2@LinApp{} = WClosed (linLayoutEnv segs lin2)
 sortLinLayouts _ LinApp{} LinEnd = error "Segments in source that do not occur in target"
 
 autoWeak :: forall segments env1 env2.
-            SSegments segments -> Layout segments env1 -> Layout segments env2 -> env1 :> env2
-autoWeak segs ly1 ly2 = sortLinLayouts segs (lineariseLayout ly1) (lineariseLayout ly2)
+            SSegments segments -> Layout True segments env1 -> Layout False segments env2 -> env1 :> env2
+autoWeak segs ly1 ly2 =
+  preWeaken segs (lineariseLayout ly1) $ \wPreweak lin1 ->
+    sortLinLayouts segs lin1 (lineariseLayout ly2) .> wPreweak