From d222c12b9c2da38ceeb531cd4b2986ff9707e018 Mon Sep 17 00:00:00 2001
From: frog-da <84839431+frog-da@users.noreply.github.com>
Date: Wed, 30 Oct 2024 22:43:07 +0300
Subject: [PATCH] Add AST and NameMap helpers

---
 .../src/Control/Monad/Foil/Internal.hs        |  7 +++++-
 .../free-foil/src/Control/Monad/Free/Foil.hs  | 22 +++++++++++++++++++
 2 files changed, 28 insertions(+), 1 deletion(-)

diff --git a/haskell/free-foil/src/Control/Monad/Foil/Internal.hs b/haskell/free-foil/src/Control/Monad/Foil/Internal.hs
index 0378c4a4..67c7b2a1 100644
--- a/haskell/free-foil/src/Control/Monad/Foil/Internal.hs
+++ b/haskell/free-foil/src/Control/Monad/Foil/Internal.hs
@@ -1,6 +1,7 @@
 {-# LANGUAGE BlockArguments             #-}
 {-# LANGUAGE ConstraintKinds            #-}
 {-# LANGUAGE DataKinds                  #-}
+{-# LANGUAGE DeriveTraversable          #-}
 {-# LANGUAGE EmptyCase                  #-}
 {-# LANGUAGE FlexibleInstances          #-}
 {-# LANGUAGE GADTs                      #-}
@@ -808,7 +809,7 @@ instance (Sinkable e) => Sinkable (Substitution e i) where
 -- * 'Name' maps
 
 -- | A /total/ map from names in scope @n@ to elements of type @a@.
-newtype NameMap (n :: S) a = NameMap { getNameMap :: IntMap a }
+newtype NameMap (n :: S) a = NameMap { getNameMap :: IntMap a } deriving (Functor, Foldable, Traversable)
 
 -- | An empty map belongs in the empty scope.
 emptyNameMap :: NameMap VoidS a
@@ -818,6 +819,10 @@ emptyNameMap = NameMap IntMap.empty
 nameMapToSubstitution :: NameMap i (e o) -> Substitution e i o
 nameMapToSubstitution (NameMap m) = (UnsafeSubstitution m)
 
+-- | Convert a 'NameMap' of expressions into a 'Scope'.
+nameMapToScope :: NameMap n a -> Scope n
+nameMapToScope (NameMap m) = UnsafeScope (IntMap.keysSet m)
+
 -- | Extend a map with multiple mappings (by repeatedly applying 'addNameBinder').
 --
 -- Note that the input list is expected to have __at least__ the same number of elements
diff --git a/haskell/free-foil/src/Control/Monad/Free/Foil.hs b/haskell/free-foil/src/Control/Monad/Free/Foil.hs
index 6f768cbd..eeb77ddd 100644
--- a/haskell/free-foil/src/Control/Monad/Free/Foil.hs
+++ b/haskell/free-foil/src/Control/Monad/Free/Foil.hs
@@ -29,8 +29,10 @@ import Data.ZipMatchK
 import           Data.Coerce                 (coerce)
 import           Data.Map                    (Map)
 import qualified Data.Map                    as Map
+import           Data.Maybe                  (mapMaybe)
 import           Data.Monoid                 (All (..))
 import           GHC.Generics                (Generic)
+import           Unsafe.Coerce               (unsafeCoerce)
 
 -- | Scoped term under a (single) name binder.
 data ScopedAST binder sig n where
@@ -376,3 +378,23 @@ convertFromScopedAST fromSig fromVar makePattern makeScoped f = \case
   ScopedAST binder body ->
     ( makePattern binder
     , makeScoped (convertFromAST fromSig fromVar makePattern makeScoped f body))
+
+-- ** Unsinking AST
+
+-- | Unsink an AST from a larger scope to a smaller scope.
+unsinkAST :: (Foil.Distinct l, Foil.CoSinkable binder, Bifoldable sig) => Foil.Scope n -> AST binder sig l -> Maybe (AST binder sig n)
+unsinkAST scope term
+  | all (`Foil.member` scope) (freeVarsOf term) = Just (unsafeCoerce term)
+  | otherwise = Nothing
+
+-- | Get the free variables of an AST.
+freeVarsOf :: (Foil.Distinct n, Foil.CoSinkable binder, Bifoldable sig) => AST binder sig n -> [Foil.Name n]
+freeVarsOf = \case
+  Var name -> [name]
+  Node node -> bifoldMap freeVarsOfScopedAST freeVarsOf node
+
+-- | Get the free variables of a scoped AST.
+freeVarsOfScopedAST :: (Foil.Distinct n, Foil.CoSinkable binder, Bifoldable sig) => ScopedAST binder sig n -> [Foil.Name n]
+freeVarsOfScopedAST (ScopedAST binder body) =
+  case Foil.assertDistinct binder of
+    Foil.Distinct -> mapMaybe (Foil.unsinkNamePattern binder) (freeVarsOf body)