ddc-core-simpl-0.4.3.1: DDC/Core/Transform/Rewrite/Match.hs
-- | Create substitution to make (subst template) == target
module DDC.Core.Transform.Rewrite.Match
( -- * Substitutions
SubstInfo
, emptySubstInfo
-- * Matching
, match)
where
import DDC.Core.Exp
import Data.Set (Set)
import Data.Map (Map)
import qualified DDC.Type.Sum as Sum
import qualified DDC.Type.Transform.AnonymizeT as T
import qualified DDC.Core.Transform.AnonymizeX as T
import qualified DDC.Core.Transform.Reannotate as T
import qualified DDC.Type.Exp.Simple as T
import qualified Data.Map as Map
import qualified Data.Set as Set
import qualified DDC.Core.Env.EnvT as EnvT
-------------------------------------------------------------------------------
-- | Value and type substition.
type SubstInfo a n
= (Map n (Exp a n), Map n (Type n))
-- | An empty substition info.
emptySubstInfo :: SubstInfo a n
emptySubstInfo
= (Map.empty, Map.empty)
lookupx n (xs,_)
= Map.lookup n xs
insertx n x (xs,tys)
= (Map.insert n x xs, tys)
-- Match Exp ------------------------------------------------------------------
-- | Create substitution to make (subst template) == target
-- Does not handle higher-order templates (ie ones with binders)
--
-- @ match emptySubstInfo (Set.fromList [r1, r2, s])
-- (stream [r1] (unstream [r2] s))
-- (stream [R0#] (unstream [R1#] (someStream 23))
--
-- => { r1 |-> R0#, r2 |-> R1, s |-> someStream 23 }
-- @
match :: (Show a, Show n, Ord n)
=> SubstInfo a n -- ^ Current substitution
-> Set n -- ^ Variables we're interested in
-> Exp a n -- ^ Template expression.
-> Exp a n -- ^ Target expression.
-> Maybe (SubstInfo a n)
-- Variables bound by the rule: restricted to just UName earlier.
match m bs (XVar _ (UName n)) r
| n `Set.member` bs
-- Check if it's already been matched
= case lookupx n m of
Nothing -> return $ insertx n r m
Just x
-> -- Check if they're equal. Anonymize so names don't matter.
-- Reannotate so annotations are ignored.
let x' = T.anonymizeX $ T.reannotate (const ()) x
r' = T.anonymizeX $ T.reannotate (const ()) r
in if x' == r'
then Just m
else Nothing
match m _ (XVar _ v1) (XVar _ v2)
| v1 == v2 = Just m
match m _ (XCon _ c1) (XCon _ c2)
| c1 == c2 = Just m
match m bs (XApp _ x11 x12) (XApp _ x21 x22)
= do m' <- match m bs x11 x21
match m' bs x12 x22
match m bs (XCast _ c1 x1) (XCast _ c2 x2)
| eqCast c1 c2
= match m bs x1 x2
match (xs, tys) bs (XType _ t1) (XType _ t2)
= do tys' <- matchT t1 t2 bs tys
return (xs, tys')
match m _ (XWitness _ w1) (XWitness _ w2)
| eqWit w1 w2 = return m
match _ _ _ _
= Nothing
eqCast :: Ord n => Cast a n -> Cast a n -> Bool
eqCast lc rc
= clean lc == clean rc
where clean c
= T.reannotate (const ())
$ case c of
CastWeakenEffect eff -> CastWeakenEffect $ T.anonymizeT eff
CastPurify wit -> CastPurify wit
CastBox -> CastBox
CastRun -> CastRun
eqWit :: Ord n => Witness a n -> Witness a n -> Bool
eqWit lw rw
= T.reannotate (const ()) lw
== T.reannotate (const ()) rw
-- Types ----------------------------------------------------------------------
type VarSet n = Set.Set n
type Subst n = Map.Map n (Type n)
-- | Try to find a simple substitution between two types.
-- Ignoring complicated effect sums.
-- Also ignoring TForall - checkRewriteRule outlaws foralls in the template, so it's safe.
-- Eg given template @a -> b@ and target @Int -> Float@,
-- returns substitution:
-- @{ a |-> Int, b |-> Float }@
--
matchT :: Ord n
=> Type n -- ^ Template type.
-> Type n -- ^ Target type.
-> VarSet n -- ^ Only attempt to match these names.
-> Subst n -- ^ Already matched (or @Map.empty@)
-> Maybe (Subst n)
matchT t1 t2 vs subst
= let t1' = unpackSumT $ T.crushSomeT EnvT.empty t1
t2' = unpackSumT $ T.crushSomeT EnvT.empty t2
in case (t1', t2') of
-- Constructor names must be equal.
--
-- Will this still work when it's a TyConBound - basically same as TVar?
(TCon tc1, TCon tc2)
| tc1 == tc2
-> Just subst
-- Decend into applications.
(TApp t11 t12, TApp t21 t22)
-> matchT t11 t21 vs subst >>= matchT t12 t22 vs
-- Sums are equivalent if all of their components are.
-- Very simple matching, only consider equivalent if both have same
-- length and in the same order.
--
-- > (Read + Write + a) `matchT` (Read + Write + Alloc)
-- > =
-- > Just { a |-> Alloc }
-- but
-- > (Read + a + Write) `matchT` (Read + Write + Alloc)
-- > =
-- > Nothing
-- and
-- > (Read + Write + Alloc + a) `matchT` (Read + Write + Alloc)
-- > =
-- > Nothing
-- despite a valid substitution existing as
-- > { a |-> !0 }
--
(TSum ts1, TSum ts2)
-> let ts1' = Sum.toList ts1
ts2' = Sum.toList ts2
go (l:ls) (r:rs) s = matchT l r vs s >>= go ls rs
go _ _ s = Just s
in if length ts1' /= length ts2'
then Nothing
else go ts1' ts2' subst
-- If template is in variable set, push the target into substitution
(TVar (UName n), _)
| Set.member n vs
, Nothing <- Map.lookup n subst
-> Just $ Map.insert n t2' subst
| Set.member n vs
, Just t1'' <- Map.lookup n subst
, T.equivT EnvT.empty t1'' t2'
-> Just subst
-- Both are variables but it's not a template variable,
-- so it's only valid if they're equal.
(TVar (UName n), TVar v2)
| not $ Set.member n vs
, UName n == v2
-> Just subst
(TVar (UIx i), TVar v2)
| UIx i == v2
-> Just subst
(TVar (UPrim n t), TVar v2)
| UPrim n t == v2
-> Just subst
-- Otherwise the two are different
(_, _) -> Nothing
-- | Unpack single element sums into plain types.
unpackSumT :: Type n -> Type n
unpackSumT (TSum ts)
| [t] <- Sum.toList ts = t
unpackSumT tt = tt