content: -- | Create substitution to make (subst template) == target
  module DDC.Core.Transform.Rewrite.Match
          ( -- * Substitutions
            SubstInfo
          , emptySubstInfo
  
            -- * Matching
          , match
          , matchArg)
  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
          matchArg m' bs x12 x22
  
  match m bs (XCast _ c1 x1) (XCast _ c2 x2)
   | eqCast c1 c2
   = match m bs x1 x2
  
  match _ _ _ _
   = Nothing
  
  
  matchArg (xs, tys) bs (RType t1) (RType t2)
   = do   tys'    <- matchT t1 t2 bs tys
          return (xs, tys')
  
  matchArg m bs         (RTerm x1) (RTerm x2)
   = do   match m bs x1 x2
  
  matchArg m _          (RWitness w1) (RWitness w2)
   | eqWit w1 w2  = return m
  
  matchArg _ _ _ _
   = 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), TVar v2)
  --          | UPrim n == 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
  
  