aboutsummaryrefslogtreecommitdiff
path: root/src/Haskell/Rewrite.hs
diff options
context:
space:
mode:
Diffstat (limited to 'src/Haskell/Rewrite.hs')
-rw-r--r--src/Haskell/Rewrite.hs95
1 files changed, 80 insertions, 15 deletions
diff --git a/src/Haskell/Rewrite.hs b/src/Haskell/Rewrite.hs
index 498ccae..2e3f6c5 100644
--- a/src/Haskell/Rewrite.hs
+++ b/src/Haskell/Rewrite.hs
@@ -1,3 +1,4 @@
+{-# LANGUAGE TupleSections #-}
module Haskell.Rewrite
(rewrite
,betared, etared, casered
@@ -9,27 +10,91 @@ import Control.Monad
import Data.List
import Data.Maybe
import qualified Data.Map.Strict as Map
+import qualified Data.Set as Set
+
import Haskell.AST
import Util
rewrite :: Name -> Expr -> Expr -> Expr
-rewrite name repl = \case
- App e as -> App (rewrite name repl e) (map (rewrite name repl) as)
- Ref n -> if n == name then repl else Ref n
- Num k -> Num k
- Tup es -> Tup (map (rewrite name repl) es)
- Lam ns e -> if name `elem` ns then Lam ns e else Lam ns (rewrite name repl e)
- Case e as -> Case (rewrite name repl e) (map caseArm as)
+rewrite target repl topexpr = fst (rewrite' mempty topexpr)
where
- caseArm (p, e') =
- if name `elem` boundVars p then (p, e') else (p, rewrite name repl e')
-
-boundVars :: Pat -> [Name]
-boundVars PatAny = []
-boundVars (PatVar n) = [n]
-boundVars (PatCon n ps) = nub $ [n] ++ concatMap boundVars ps
-boundVars (PatTup ps) = nub $ concatMap boundVars ps
+ -- When moving into the subexpression E under a binding for some variable
+ -- 'x', if 'x' is free in 'repl', we need to alpha-rename 'x' to something
+ -- that is:
+ -- - Not free in 'repl';
+ -- - Not free in E;
+ -- - Not equal to a name 'y' bound in E if the subexpression under the
+ -- binder for 'y' has an occurrence of 'x'.
+ -- If we strengthen the third requirement to prohibit all bound variables
+ -- in E, the second and third together mean "all variables in E". This is
+ -- what we will use.
+ frees :: Set.Set Name
+ frees = freeVariables repl
+
+ -- Returns rewritten expression, and whether any actual rewrites were
+ -- performed. This allows preventing unnecessary alpha-renames.
+ rewrite' :: Map.Map Name Name -> Expr -> (Expr, Bool)
+ rewrite' mapping = \case
+ App e as -> let (e' : as', bs) = unzip (map (rewrite' mapping) (e : as))
+ in (App e' as', or bs)
+ Ref n -> case Map.lookup n mapping of
+ Just n' -> (Ref n', False) -- renamed variable cannot be target
+ Nothing | n == target -> (repl, True)
+ | otherwise -> (Ref n, False)
+ Con n -> (Con n, False)
+ Num k -> (Num k, False)
+ Tup es -> let (es', bs) = unzip (map (rewrite' mapping) es)
+ in (Tup es', or bs)
+ Lam ns e
+ | target `elem` ns -> (Lam ns e, False)
+ | otherwise ->
+ let forbidden = frees <> allVars e <> Set.fromList ns
+ -- Note that Map.<> is left-preferring
+ mapping' = freshenL ns frees forbidden <> mapping
+ ns' = map (rename mapping') ns
+ (e', b) = rewrite' mapping' e
+ in if b then (Lam ns' e', True) else (Lam ns e, False)
+ Case e as ->
+ let (scrutinee, b1) = rewrite' mapping e
+ (arms, bs) =
+ unzip [if target `elem` ns
+ then ((p, e'), False)
+ else let forbidden = frees <> allVars e' <> Set.fromList ns
+ mapping' = freshenL ns frees forbidden <> mapping
+ (rhs, b) = rewrite' mapping' e'
+ in if b
+ then ((renamePat mapping' p, rhs), True)
+ else ((p, e'), False)
+ | (p, e') <- as
+ , let ns = Set.toList (boundVars p)]
+ in (Case scrutinee arms, b1 || or bs)
+
+ -- Freshens all variables found in 'vars' to something that is not in
+ -- 'bnd', returning a replace mapping. Later names override earlier ones
+ -- in the mapping.
+ freshenL ns vars bnd =
+ foldl (\mp n ->
+ if n `Set.member` vars
+ then Map.insert n (freshName bnd n) mp
+ else mp)
+ mempty ns
+
+ -- Finds a fresh name for 'n' that is not in 'bnd'.
+ freshName bnd n =
+ head [n ++ "_R_" ++ show i
+ | i <- [1::Int ..]
+ , let n' = n ++ show i
+ , n' `Set.notMember` bnd]
+
+ renamePat mapping = \case
+ PatAny -> PatAny
+ PatVar n -> PatVar (rename mapping n)
+ PatCon n ps -> PatCon (rename mapping n) (map (renamePat mapping) ps)
+ PatTup ps -> PatTup (map (renamePat mapping) ps)
+
+ rename :: Map.Map Name Name -> Name -> Name
+ rename mapping n = fromMaybe n (Map.lookup n mapping)
betared :: Expr -> Expr
betared (App (Lam (n:as) bd) (arg:args)) =