Generic alpha-equivalence

emilaxelsson · Jun 16, 2015 · b11412a · b11412a
1 parent a7e6f20
commit b11412a
Show file tree

Hide file tree

Showing 3 changed files with 80 additions and 23 deletions.
diff --git a/examples/Feldspar.hs b/examples/Feldspar.hs
@@ -95,6 +95,23 @@ instance HasVars ExpF Name
           ([],[v']) -> Arr l v' f
     renameVars f = fmap fst f
 
+instance EqConstr ExpF
+  where
+    eqConstr (Var v1)     (Var v2)     = v1==v2
+    eqConstr (LitB b1)    (LitB b2)    = b1==b2
+    eqConstr (LitI i1)    (LitI i2)    = i1==i2
+    eqConstr (Add _ _)    (Add _ _)    = True
+    eqConstr (Sub _ _)    (Sub _ _)    = True
+    eqConstr (Mul _ _)    (Mul _ _)    = True
+    eqConstr (Eq _ _)     (Eq _ _)     = True
+    eqConstr (Min _ _)    (Min _ _)    = True
+    eqConstr (Max _ _)    (Max _ _)    = True
+    eqConstr (If _ _ _)   (If _ _ _)   = True
+    eqConstr (Let v1 _ _) (Let v2 _ _) = v1==v2
+    eqConstr (Arr _ v1 _) (Arr _ v2 _) = v1==v2
+    eqConstr (Ix _ _)     (Ix _ _)     = True
+    eqConstr _ _ = False
+
 -- | The type of Feldspar expressions
 newtype Data a = Data { unData :: Tree ExpF }
 
@@ -330,6 +347,9 @@ rangeMax (l1,u1) (l2,u2) = (max l1 l2, max u1 u2)
 rename' :: Dag ExpF -> Dag ExpF
 rename' = rename (Just (0 :: Name))
 
+alphaEq' :: Tree ExpF -> Tree ExpF -> Bool
+alphaEq' = alphaEq (Nothing :: Maybe Name)
+
 -- | Size of an expression = over-approximation of the set of values it might
 -- take on
 --

diff --git a/examples/Paper.hs b/examples/Paper.hs
@@ -408,30 +408,23 @@ instance HasVars ExpF Name
           ([],[v']) -> Iter' v' k i b
     renameVars f = fmap fst f
 
+instance EqConstr ExpF
+  where
+    eqConstr (LitB' b1)      (LitB' b2)      = b1==b2
+    eqConstr (LitI' i1)      (LitI' i2)      = i1==i2
+    eqConstr (Var' v1)       (Var' v2)       = v1==v2
+    eqConstr (Eq' _ _)       (Eq' _ _)       = True
+    eqConstr (Add' _ _)      (Add' _ _)      = True
+    eqConstr (If' _ _ _)     (If' _ _ _)     = True
+    eqConstr (Iter' _ _ _ _) (Iter' _ _ _ _) = True
+    eqConstr _ _ = False
+
 -- | Make the DAG well-scoped
 rename' :: Dag ExpF -> Dag ExpF
 rename' = rename (Just "")
 
-alphaEq' :: [(Name,Name)] -> Tree ExpF -> Tree ExpF -> Bool
-alphaEq' env (In (Var' v1)) (In (Var' v2)) =
-    case (lookup v1 env, lookup v2 env') of
-      (Nothing, Nothing)   -> v1==v2  -- Free variables
-      (Just v2', Just v1') -> v1==v1' && v2==v2'
-      _                    -> False
-  where
-    env' = [(v2,v1) | (v1,v2) <- env]
-alphaEq' env (In (Iter' v1 k1 i1 b1)) (In (Iter' v2 k2 i2 b2)) =
-    alphaEq' env k1 k2 && alphaEq' env i1 i2 && alphaEq' ((v1,v2):env) b1 b2
-alphaEq' env (In (LitB' b1))     (In (LitB' b2))     = b1==b2
-alphaEq' env (In (LitI' i1))     (In (LitI' i2))     = i1==i2
-alphaEq' env (In (Eq' a1 b1))    (In (Eq' a2 b2))    = alphaEq' env a1 a2 && alphaEq' env b1 b2
-alphaEq' env (In (Add' a1 b1))   (In (Add' a2 b2))   = alphaEq' env a1 a2 && alphaEq' env b1 b2
-alphaEq' env (In (If' c1 t1 f1)) (In (If' c2 t2 f2)) = alphaEq' env c1 c2 && alphaEq' env t1 t2 && alphaEq' env f1 f2
-alphaEq' env _ _ = False
-
--- | Alpha-equivalence (for testing the renamer)
-alphaEq :: Tree ExpF -> Tree ExpF -> Bool
-alphaEq = alphaEq' []
+alphaEq' :: Tree ExpF -> Tree ExpF -> Bool
+alphaEq' = alphaEq (Nothing :: Maybe Name)
 
 -- | Like 'flatten' but adds the root as a node in the graph
 flatten' :: Traversable f => Dag f -> (Node, IntMap (f Node), Int)
@@ -472,7 +465,7 @@ isWellScoped g = all checkVar $ fmap concat $ groups sc
     sc = scope g []
 
 -- | Renaming does not changes semantics
-prop_rename1 g = unravelDag g `alphaEq` unravelDag (rename' g)
+prop_rename1 g = unravelDag g `alphaEq'` unravelDag (rename' g)
 
 -- | Renaming does not decrease the number of edges
 prop_rename2 g = length (IntMap.toList $ edges g) <= length (IntMap.toList $ edges $ rename' g)

diff --git a/src/Variables.hs b/src/Variables.hs
@@ -1,13 +1,22 @@
 {-# LANGUAGE MultiParamTypeClasses #-}
+{-# LANGUAGE ScopedTypeVariables #-}
 
 -- | A generic interface to constructs that bind or represent variables
 
-module Variables where
+module Variables
+    ( IsVar (..)
+    , HasVars (..)
+    , EqConstr (..)
+    , alphaEq
+    ) where
 
 
 
-import Data.Set
+import qualified Data.Foldable as Foldable
+import Data.Set (Set)
+import qualified Data.Set as Set
 
+import Tree
 import Mapping
 
 
@@ -33,3 +42,38 @@ class (Functor f, IsVar v) => HasVars f v where
     renameVars :: f (a, Set v) -> f a
     renameVars f = fmap fst f
 
+class EqConstr f where
+    eqConstr :: f a -> f a -> Bool
+
+alphaEq' :: forall v f . (EqConstr f, HasVars f v, Traversable f, Eq v)
+    => [(v,v)] -> Tree f -> Tree f -> Bool
+alphaEq' env (In var1) (In var2)
+    | Just v1 <- isVar var1
+    , Just v2 <- isVar var2
+    = case (lookup v1 env, lookup v2 env') of
+        (Nothing, Nothing)   -> v1==v2  -- Free variables
+        (Just v2', Just v1') -> v1==v1' && v2==v2'
+        _                    -> False
+  where
+    env' = [(v2,v1) | (v1,v2) <- env]
+alphaEq' env (In f) (In g)
+    =  eqConstr f g
+    && all (checkChildren env)
+        ( zip
+            (Foldable.toList $ number f)
+            (Foldable.toList $ number g)
+        )
+  where
+    checkChildren :: [(v,v)] -> (Numbered (Tree f), Numbered (Tree f)) -> Bool
+    checkChildren env (Numbered i1 t1, Numbered i2 t2)
+        = length vs1 == length vs2 && alphaEq' (zip vs1 vs2 ++ env) t1 t2
+      where
+        vs1 = Set.toList $ lookupNumMap Set.empty i1 (bindsVars $ number f)
+        vs2 = Set.toList $ lookupNumMap Set.empty i2 (bindsVars $ number g)
+
+-- | Alpha-equivalence
+alphaEq :: forall proxy v f
+    .  (EqConstr f, HasVars f v, Traversable f, Eq v)
+    => proxy v -> Tree f -> Tree f -> Bool
+alphaEq _ = alphaEq' ([] :: [(v,v)])
+