Generalize extendScopePattern

fizruk · Jun 20, 2024 · 1424051 · 1424051
1 parent 36fbb27
commit 1424051
Show file tree

Hide file tree

Showing 3 changed files with 28 additions and 8 deletions.
diff --git a/haskell/free-foil/src/Control/Monad/Foil.hs b/haskell/free-foil/src/Control/Monad/Foil.hs
@@ -21,6 +21,7 @@ module Control.Monad.Foil (
   NameBinder,
   emptyScope,
   extendScope,
+  extendScopePattern,
   member,
   nameOf,
   nameId,

diff --git a/haskell/free-foil/src/Control/Monad/Foil/Internal.hs b/haskell/free-foil/src/Control/Monad/Foil/Internal.hs
@@ -76,6 +76,33 @@ extendScope :: NameBinder n l -> Scope n -> Scope l
 extendScope (UnsafeNameBinder (UnsafeName name)) (UnsafeScope scope) =
   UnsafeScope (IntSet.insert name scope)
 
+newtype ExtendScope n l (o :: S) (o' :: S) = ExtendScope (Scope n -> Scope l)
+
+idExtendScope :: ExtendScope n n o o'
+idExtendScope = ExtendScope id
+
+compExtendScope
+  :: ExtendScope n i o o'
+  -> ExtendScope i l o' o''
+  -> ExtendScope n l o o''
+compExtendScope (ExtendScope f) (ExtendScope g)
+  = ExtendScope (g . f)
+
+-- | Extend scope with variables inside a pattern.
+-- This is a more flexible version of 'extendScope'.
+extendScopePattern
+  :: (Distinct n, CoSinkable pattern)
+  => pattern n l -> Scope n -> Scope l
+extendScopePattern pat scope = withPattern
+  (\_scope' binder k ->
+    unsafeAssertFresh binder $ \binder' ->
+      k (ExtendScope (extendScope binder)) binder')
+  idExtendScope
+  compExtendScope
+  scope
+  pat
+  (\(ExtendScope extend) _ -> extend scope)
+
 -- | A runtime check for potential name capture.
 member :: Name l -> Scope n -> Bool
 member (UnsafeName name) (UnsafeScope s) = rawMember name s

diff --git a/haskell/lambda-pi/src/Language/LambdaPi/Impl/Foil.hs b/haskell/lambda-pi/src/Language/LambdaPi/Impl/Foil.hs
@@ -230,14 +230,6 @@ namesOfPattern (PatternPair l r) =
     (Ext, Distinct) -> case (assertExtPattern r, assertDistinctPattern r) of
        (Ext, Distinct) -> map sink (namesOfPattern l) ++ namesOfPattern r
 
--- | Extend scope with variables inside a pattern.
--- This is a more flexible version of 'extendScope'.
-extendScopePattern :: Pattern n l -> Scope n -> Scope l
-extendScopePattern = \case
-  PatternWildcard -> id
-  PatternVar binder -> extendScope binder
-  PatternPair l r -> extendScopePattern r . extendScopePattern l
-
 -- | Perform substitution in a \(\lambda\Pi\)-term.
 substitute :: Distinct o => Scope o -> Substitution Expr i o -> Expr i -> Expr o
 substitute scope subst = \case