Back recursor types with type families

src/Juvix/Compiler/Core/Extra.hs

@@ -4,6 +4,8 @@ module Juvix.Compiler.Core.Extra
     module Juvix.Compiler.Core.Extra.Recursors,
     module Juvix.Compiler.Core.Extra.Info,
     module Juvix.Compiler.Core.Extra.Equality,
+    module Juvix.Compiler.Core.Extra.Recursors.Fold.Named,
+    module Juvix.Compiler.Core.Extra.Recursors.Map.Named,
   )
 where
 
@@ -12,30 +14,32 @@ import Juvix.Compiler.Core.Extra.Base
 import Juvix.Compiler.Core.Extra.Equality
 import Juvix.Compiler.Core.Extra.Info
 import Juvix.Compiler.Core.Extra.Recursors
+import Juvix.Compiler.Core.Extra.Recursors.Fold.Named
+import Juvix.Compiler.Core.Extra.Recursors.Map.Named
 import Juvix.Compiler.Core.Language
 
 isClosed :: Node -> Bool
 isClosed = not . has freeVars
 
-getFreeVars :: Node -> HashSet Index
+getFreeVars :: Node -> HashSet Var
 getFreeVars n = HashSet.fromList (n ^.. freeVars)
 
-freeVars :: SimpleFold Node Index
-freeVars f = ufoldAN' reassemble go
+freeVars :: SimpleFold Node Var
+freeVars f = ufoldNA reassemble go
   where
     go k = \case
-      NVar (Var i idx)
-        | idx >= k -> mkVar i <$> f (idx - k)
+      NVar var@(Var _ idx)
+        | idx >= k -> NVar <$> f var
       n -> pure n
 
-getIdents :: Node -> HashSet Symbol
+getIdents :: Node -> HashSet Ident
 getIdents n = HashSet.fromList (n ^.. nodeIdents)
 
-nodeIdents :: Traversal' Node Symbol
-nodeIdents f = umapLeaves go
+nodeIdents :: Traversal' Node Ident
+nodeIdents f = ufoldA reassemble go
   where
     go = \case
-      NIdt (Ident i d) -> mkIdent i <$> f d
+      NIdt i -> NIdt <$> f i
       n -> pure n
 
 countFreeVarOccurrences :: Index -> Node -> Int

src/Juvix/Compiler/Core/Extra/Recursors.hs

@@ -1,49 +1,17 @@
 module Juvix.Compiler.Core.Extra.Recursors
-  ( module Juvix.Compiler.Core.Extra.Recursors,
-    module Juvix.Compiler.Core.Extra.Recursors.Applicative,
+  ( module Juvix.Compiler.Core.Extra.Recursors.Fold,
     module Juvix.Compiler.Core.Extra.Recursors.Collector,
-    module Juvix.Compiler.Core.Extra.Recursors.Monadic,
+    module Juvix.Compiler.Core.Extra.Recursors.Map,
+    module Juvix.Compiler.Core.Extra.Recursors.Map.Named,
+    module Juvix.Compiler.Core.Extra.Recursors.Fold.Named,
+    module Juvix.Compiler.Core.Extra.Recursors.Recur,
   )
 where
 
-import Data.Functor.Identity
-import Juvix.Compiler.Core.Extra.Recursors.Applicative
 import Juvix.Compiler.Core.Extra.Recursors.Base
 import Juvix.Compiler.Core.Extra.Recursors.Collector
-import Juvix.Compiler.Core.Extra.Recursors.Monadic
-
-dmap :: (Node -> Node) -> Node -> Node
-dmap f = runIdentity . dmapM (return . f)
-
-dmapB :: (BinderList Info -> Node -> Node) -> Node -> Node
-dmapB f = runIdentity . dmapMB (\is -> return . f is)
-
-dmapN :: (Index -> Node -> Node) -> Node -> Node
-dmapN f = runIdentity . dmapMN (\idx -> return . f idx)
-
-ufold :: (a -> a -> a) -> (Node -> a) -> Node -> a
-ufold uplus f n = runIdentity $ ufoldA uplus (return . f) n
-
-ufoldB :: (a -> a -> a) -> (BinderList Info -> Node -> a) -> Node -> a
-ufoldB uplus f = runIdentity . ufoldAB uplus (\is -> return . f is)
-
-ufoldN :: (a -> a -> a) -> (Index -> Node -> a) -> Node -> a
-ufoldN uplus f = runIdentity . ufoldAN uplus (\idx -> return . f idx)
-
-gather :: (a -> Node -> a) -> a -> Node -> a
-gather f acc = run . execState acc . walk (\n' -> modify' (`f` n'))
-
-gatherB :: (BinderList Info -> a -> Node -> a) -> a -> Node -> a
-gatherB f acc = run . execState acc . walkB (\is n' -> modify' (\a -> f is a n'))
-
-gatherN :: (Index -> a -> Node -> a) -> a -> Node -> a
-gatherN f acc = run . execState acc . walkN (\idx n' -> modify' (\a -> f idx a n'))
-
-umap :: (Node -> Node) -> Node -> Node
-umap f = runIdentity . umapM (return . f)
-
-umapB :: (BinderList Info -> Node -> Node) -> Node -> Node
-umapB f = runIdentity . umapMB (\is -> return . f is)
-
-umapN :: (Index -> Node -> Node) -> Node -> Node
-umapN f = runIdentity . umapMN (\idx -> return . f idx)
+import Juvix.Compiler.Core.Extra.Recursors.Fold
+import Juvix.Compiler.Core.Extra.Recursors.Fold.Named
+import Juvix.Compiler.Core.Extra.Recursors.Map
+import Juvix.Compiler.Core.Extra.Recursors.Map.Named
+import Juvix.Compiler.Core.Extra.Recursors.Recur

src/Juvix/Compiler/Core/Extra/Recursors/Base.hs

@@ -3,10 +3,12 @@ module Juvix.Compiler.Core.Extra.Recursors.Base
     module Juvix.Compiler.Core.Language,
     module Juvix.Compiler.Core.Info.BinderInfo,
     module Juvix.Compiler.Core.Extra.Recursors.Collector,
+    module Juvix.Compiler.Core.Extra.Recursors.Recur,
   )
 where
 
 import Juvix.Compiler.Core.Data.BinderList (BinderList)
 import Juvix.Compiler.Core.Extra.Recursors.Collector
+import Juvix.Compiler.Core.Extra.Recursors.Recur
 import Juvix.Compiler.Core.Info.BinderInfo
 import Juvix.Compiler.Core.Language

src/Juvix/Compiler/Core/Extra/Recursors/Collector.hs

@@ -16,11 +16,19 @@ makeLenses ''Collector
 unitCollector :: Collector a ()
 unitCollector = Collector () (\_ _ -> ())
 
+binderInfoCollector' :: BinderList Info -> Collector (Int, [Info]) (BinderList Info)
+binderInfoCollector' ini = Collector ini collect
+  where
+    collect :: (Int, [Info]) -> BinderList Info -> BinderList Info
+    collect (k, bi) c
+      | k == 0 = c
+      | otherwise = BL.prepend (reverse bi) c
+
 binderInfoCollector :: Collector (Int, [Info]) (BinderList Info)
-binderInfoCollector =
-  Collector
-    mempty
-    (\(k, bi) c -> if k == 0 then c else BL.prepend (reverse bi) c)
+binderInfoCollector = binderInfoCollector' mempty
+
+binderNumCollector' :: Int -> Collector (Int, [Info]) Index
+binderNumCollector' ini = Collector ini (\(k, _) c -> c + k)
 
 binderNumCollector :: Collector (Int, [Info]) Index
-binderNumCollector = Collector 0 (\(k, _) c -> c + k)
+binderNumCollector = binderNumCollector' 0

src/Juvix/Compiler/Core/Extra/Recursors/Fold.hs

@@ -0,0 +1,31 @@
+-- | Fold recursors over 'Node'.
+module Juvix.Compiler.Core.Extra.Recursors.Fold where
+
+import Juvix.Compiler.Core.Extra.Base
+import Juvix.Compiler.Core.Extra.Recursors.Base
+
+ufoldG ::
+  forall c a f.
+  Applicative f =>
+  Collector (Int, [Info]) c ->
+  (a -> [a] -> a) ->
+  (c -> Node -> f a) ->
+  Node ->
+  f a
+ufoldG coll uplus f = go (coll ^. cEmpty)
+  where
+    go :: c -> Node -> f a
+    go c n = do
+      mas' <- sequenceA mas
+      n' <- f c n
+      pure (uplus n' mas')
+      where
+        ni :: NodeDetails
+        ni = destruct n
+        mas :: [f a]
+        mas =
+          zipWith3Exact
+            (\n' k bis -> go ((coll ^. cCollect) (k, bis) c) n')
+            (ni ^. nodeChildren)
+            (ni ^. nodeChildBindersNum)
+            (ni ^. nodeChildBindersInfo)

src/Juvix/Compiler/Core/Extra/Recursors/Fold/Named.hs

@@ -0,0 +1,41 @@
+module Juvix.Compiler.Core.Extra.Recursors.Fold.Named where
+
+import Data.Functor.Identity
+import Juvix.Compiler.Core.Extra.Recursors.Base
+import Juvix.Compiler.Core.Extra.Recursors.Fold
+
+ufoldA :: Applicative f => (a -> [a] -> a) -> (Node -> f a) -> Node -> f a
+ufoldA uplus f = ufoldG unitCollector uplus (const f)
+
+ufoldLA :: Applicative f => (a -> [a] -> a) -> (BinderList Info -> Node -> f a) -> Node -> f a
+ufoldLA uplus f = ufoldG binderInfoCollector uplus f
+
+ufoldNA :: Applicative f => (a -> [a] -> a) -> (Index -> Node -> f a) -> Node -> f a
+ufoldNA uplus f = ufoldG binderNumCollector uplus f
+
+walk :: Applicative f => (Node -> f ()) -> Node -> f ()
+walk = ufoldA (foldr mappend)
+
+walkN :: Applicative f => (Index -> Node -> f ()) -> Node -> f ()
+walkN = ufoldNA (foldr mappend)
+
+walkL :: Applicative f => (BinderList Info -> Node -> f ()) -> Node -> f ()
+walkL = ufoldLA (foldr mappend)
+
+ufold :: (a -> [a] -> a) -> (Node -> a) -> Node -> a
+ufold uplus f = runIdentity . ufoldA uplus (return . f)
+
+ufoldL :: (a -> [a] -> a) -> (BinderList Info -> Node -> a) -> Node -> a
+ufoldL uplus f = runIdentity . ufoldLA uplus (\is -> return . f is)
+
+ufoldN :: (a -> [a] -> a) -> (Index -> Node -> a) -> Node -> a
+ufoldN uplus f = runIdentity . ufoldNA uplus (\idx -> return . f idx)
+
+gather :: (a -> Node -> a) -> a -> Node -> a
+gather f acc = run . execState acc . walk (\n' -> modify' (`f` n'))
+
+gatherL :: (BinderList Info -> a -> Node -> a) -> a -> Node -> a
+gatherL f acc = run . execState acc . walkL (\is n' -> modify' (\a -> f is a n'))
+
+gatherN :: (Index -> a -> Node -> a) -> a -> Node -> a
+gatherN f acc = run . execState acc . walkN (\idx n' -> modify' (\a -> f idx a n'))