Support nested containers in TH for Signature and pattern synonyms

fizruk · Oct 22, 2024 · 45cbc2e · 45cbc2e
1 parent 85f7d16
commit 45cbc2e
Show file tree

Hide file tree

Showing 2 changed files with 74 additions and 8 deletions.
diff --git a/haskell/free-foil/src/Control/Monad/Free/Foil/TH/PatternSynonyms.hs b/haskell/free-foil/src/Control/Monad/Free/Foil/TH/PatternSynonyms.hs
@@ -69,21 +69,56 @@ mkPatternSynonym signatureType scope term = \case
       Left ((b, _), (x, _)) -> ConP 'ScopedAST [] [VarP b, VarP x]
       Right (x, _) -> VarP x
 
-    toPatternArgType i (_bang, VarT typeName)
+    toPatternArgType i (_bang, type_@(VarT typeName))
       | typeName == scope =
           Left
             ( (mkName ("b" ++ show i), foldl AppT binderT [VarT n, VarT l])
-            , (mkName ("x" ++ show i), PeelConT ''AST [binderT, signatureType, VarT l]))
+            , (mkName ("x" ++ show i), replaceScopeTermInType l type_))
       | typeName == term =
-          Right (mkName ("x" ++ show i), PeelConT ''AST [binderT, signatureType, VarT n])
+          Right (mkName ("x" ++ show i), replaceScopeTermInType l type_)
       where
         l = mkName ("l" ++ show i)
     toPatternArgType i (_bang, type_)
-      = Right (mkName ("z" ++ show i), type_)
+      = Right (mkName ("z" ++ show i), replaceScopeTermInType l type_)
+      where
+        l = mkName ("l" ++ show i)
 
     mkPatternName conName = mkName (dropEnd (length "Sig") (nameBase conName))
     dropEnd k = reverse . drop k . reverse
 
     collapse = \case
       Left (x, y) -> [x, y]
       Right x -> [x]
+
+    replaceScopeTermInType lscope = \case
+      VarT typeName
+        | typeName == scope -> PeelConT ''AST [binderT, signatureType, VarT lscope]
+        | typeName == term  -> PeelConT ''AST [binderT, signatureType, VarT n]
+      ForallT bndrs ctx type_ -> ForallT bndrs ctx (replaceScopeTermInType lscope type_)
+      ForallVisT bndrs type_ -> ForallVisT bndrs (replaceScopeTermInType lscope type_)
+      AppT f x -> AppT (replaceScopeTermInType lscope f) (replaceScopeTermInType lscope x)
+      AppKindT f k -> AppKindT (replaceScopeTermInType lscope f) k
+      SigT t k -> SigT (replaceScopeTermInType lscope t) k
+      t@ConT{} -> t
+      t@VarT{} -> t
+      t@PromotedT{} -> t
+      InfixT l op r -> InfixT (replaceScopeTermInType lscope l) op (replaceScopeTermInType lscope r)
+      UInfixT l op r -> UInfixT (replaceScopeTermInType lscope l) op (replaceScopeTermInType lscope r)
+      PromotedInfixT l op r -> PromotedInfixT (replaceScopeTermInType lscope l) op (replaceScopeTermInType lscope r)
+      PromotedUInfixT l op r -> PromotedUInfixT (replaceScopeTermInType lscope l) op (replaceScopeTermInType lscope r)
+      ParensT t -> ParensT (replaceScopeTermInType lscope t)
+      t@TupleT{} -> t
+      t@UnboxedTupleT{} -> t
+      t@UnboxedSumT{} -> t
+      t@ArrowT{} -> t
+      t@MulArrowT{} -> t
+      t@EqualityT{} -> t
+      t@ListT{} -> t
+      t@PromotedTupleT{} -> t
+      t@PromotedNilT{} -> t
+      t@PromotedConsT{} -> t
+      t@StarT{} -> t
+      t@ConstraintT{} -> t
+      t@LitT{} -> t
+      t@WildCardT{} -> t
+      ImplicitParamT s t -> ImplicitParamT s (replaceScopeTermInType lscope t)
diff --git a/haskell/free-foil/src/Control/Monad/Free/Foil/TH/Signature.hs b/haskell/free-foil/src/Control/Monad/Free/Foil/TH/Signature.hs
@@ -89,9 +89,40 @@ mkSignature termT nameT scopeT patternT = do
           where
             k (x, y) = (name, x, y)
 
-        toSignatureParam (bang_, PeelConT typeName _typeParams)
+        toSignatureParam (_bang, PeelConT typeName _typeParams)
           | typeName == nameT = fail ("variable with other stuff in constructor: " ++ show con')
           | typeName == patternT = pure Nothing -- skip binders, they will be inserted automatically with each scoped term
-          | typeName == scopeT = pure (Just (bang_, VarT scope))
-          | typeName == termT = pure (Just (bang_, VarT term))
-        toSignatureParam bt = pure (Just bt)  -- everything else remains as is
+        toSignatureParam (bang_, type_) = pure (Just (bang_, replaceScopeTermInType type_))
+
+        replaceScopeTermInType = \case
+          PeelConT typeName _typeParams
+            | typeName == scopeT -> VarT scope
+            | typeName == termT -> VarT term
+          ForallT bndrs ctx type_ -> ForallT bndrs ctx (replaceScopeTermInType type_)
+          ForallVisT bndrs type_ -> ForallVisT bndrs (replaceScopeTermInType type_)
+          AppT f x -> AppT (replaceScopeTermInType f) (replaceScopeTermInType x)
+          AppKindT f k -> AppKindT (replaceScopeTermInType f) k
+          SigT t k -> SigT (replaceScopeTermInType t) k
+          t@ConT{} -> t
+          t@VarT{} -> t
+          t@PromotedT{} -> t
+          InfixT l op r -> InfixT (replaceScopeTermInType l) op (replaceScopeTermInType r)
+          UInfixT l op r -> UInfixT (replaceScopeTermInType l) op (replaceScopeTermInType r)
+          PromotedInfixT l op r -> PromotedInfixT (replaceScopeTermInType l) op (replaceScopeTermInType r)
+          PromotedUInfixT l op r -> PromotedUInfixT (replaceScopeTermInType l) op (replaceScopeTermInType r)
+          ParensT t -> ParensT (replaceScopeTermInType t)
+          t@TupleT{} -> t
+          t@UnboxedTupleT{} -> t
+          t@UnboxedSumT{} -> t
+          t@ArrowT{} -> t
+          t@MulArrowT{} -> t
+          t@EqualityT{} -> t
+          t@ListT{} -> t
+          t@PromotedTupleT{} -> t
+          t@PromotedNilT{} -> t
+          t@PromotedConsT{} -> t
+          t@StarT{} -> t
+          t@ConstraintT{} -> t
+          t@LitT{} -> t
+          t@WildCardT{} -> t
+          ImplicitParamT s t -> ImplicitParamT s (replaceScopeTermInType t)