content: -- |
  -- Module      : Minigent.TC.Assign
  -- Copyright   : (c) Data61 2018-2019
  --                   Commonwealth Science and Research Organisation (CSIRO)
  --                   ABN 41 687 119 230
  -- License     : BSD3
  --
  -- A module for solver assignments to unification variables.
  --
  -- May be imported unqualified.
  module Minigent.TC.Assign
    ( Assign (..) 
    , substAssign
    , substsToMaps
    , substsToMapsDisjoint
    , substConstraint'
    ) where 
  
  import Minigent.Syntax
  import Minigent.Syntax.Utils
  
  import Data.Maybe (fromMaybe)
  import qualified Data.Map as M
  
  import qualified Minigent.Syntax.Utils.Row as Row
  import qualified Minigent.Syntax.Utils.Rewrite as Rewrite
  
  -- | An assignment to a unification variable.
  data Assign = TyAssign VarName Type
              | RowAssign VarName Row
              | SigilAssign VarName Sigil
              | RecParAssign VarName RecPar
              deriving (Show)
  
  -- | Apply an assignment to a unification variable to a type.
  substAssign :: Assign -> Rewrite.Rewrite Type
  substAssign (TyAssign v t) = substUV (v, t)
  substAssign (RowAssign v t) = substRowV (v, t)
  substAssign (SigilAssign v t) = substSigilV (v, t)
  substAssign (RecParAssign v1 v2) = substRecPar (v1, v2)
  
  -- Assumes: that substitutions do not substitute for variables in other given substitutions
  substsToMaps :: [Assign] -> (M.Map VarName Type, M.Map VarName RecPar, M.Map VarName Row, M.Map VarName Sigil)
  substsToMaps [] = (M.empty, M.empty, M.empty, M.empty)
  substsToMaps (TyAssign v t : substs) =
    let (munif, mrp, mrow, ms) = substsToMaps substs
        t' = substType' munif mrp mrow ms t
     in (M.insert v t' munif, mrp, mrow, ms)
  substsToMaps (RowAssign v r : substs) =
    let (munif, mrp, mrow, ms) = substsToMaps substs
        r' = substRow' munif mrp mrow ms r
     in (munif, mrp, M.insert v r' mrow, ms)
  substsToMaps (SigilAssign v s : substs) =
    let (munif, mrp, mrow, ms) = substsToMaps substs
        s' = substSigil' ms s
     in (munif, mrp, mrow, M.insert v s' ms)
  substsToMaps (RecParAssign v rp : substs) =
    let (munif, mrp, mrow, ms) = substsToMaps substs
        rp' = substRecPar' mrp rp
     in (munif, M.insert v rp' mrp, mrow, ms)
  
  -- Assumes that no variable which is substituted for occurs in the output of the substitution
  substsToMapsDisjoint :: [Assign] -> (M.Map VarName Type, M.Map VarName RecPar, M.Map VarName Row, M.Map VarName Sigil)
  substsToMapsDisjoint [] = (M.empty, M.empty, M.empty, M.empty)
  substsToMapsDisjoint (TyAssign v t : substs) =
    let (munif, mrp, mrow, ms) = substsToMapsDisjoint substs
     in (M.insert v t munif, mrp, mrow, ms)
  substsToMapsDisjoint (RowAssign v r : substs) =
    let (munif, mrp, mrow, ms) = substsToMapsDisjoint substs
     in (munif, mrp, M.insert v r mrow, ms)
  substsToMapsDisjoint (SigilAssign v s : substs) =
    let (munif, mrp, mrow, ms) = substsToMapsDisjoint substs
     in (munif, mrp, mrow, M.insert v s ms)
  substsToMapsDisjoint (RecParAssign v rp : substs) =
    let (munif, mrp, mrow, ms) = substsToMapsDisjoint substs
     in (munif, M.insert v rp mrp, mrow, ms)
  
  
  substAssign' :: M.Map VarName Type -> 
                  M.Map VarName RecPar ->
                  M.Map VarName Row ->
                  M.Map VarName Sigil ->
                  Assign -> Assign
  substAssign' munif mrp mrow ms = go
   where
    go (TyAssign     v t ) = TyAssign v $ substType' munif mrp mrow ms t
    go (RowAssign    v r ) = RowAssign v $ substRow' munif mrp mrow ms r
    go (SigilAssign  v s ) = SigilAssign v $ substSigil' ms s
    go (RecParAssign v rp) = RecParAssign v $ substRecPar' mrp rp
  
  substConstraint' :: M.Map VarName Type -> 
                      M.Map VarName RecPar ->
                      M.Map VarName Row ->
                      M.Map VarName Sigil ->
                      Constraint -> Constraint
  substConstraint' munif mrp mrow ms = go
   where
    go (c1 :&: c2)             = go c1 :&: go c2
    go (i :<=: t)              = i :<=: substTy t
    go (Share     t)           = Share $ substTy t
    go (Drop      t)           = Drop $ substTy t
    go (Escape    t)           = Escape $ substTy t
    go (Exhausted t)           = Exhausted $ substTy t
    go (t1  :<  t2 )           = substTy t1 :< substTy t2
    go (t1  :=: t2 )           = substTy t1 :=: substTy t2
    go (Solved t)              = Solved $ substTy t
    go (UnboxedNoRecurse rp s) = UnboxedNoRecurse (substRecPar' mrp rp) (substSigil' ms s)
    go Sat                     = Sat
    go Unsat                   = Unsat
  
    substTy = substType' munif mrp mrow ms
  
  substType' :: M.Map VarName Type -> 
                M.Map VarName RecPar ->
                M.Map VarName Row ->
                M.Map VarName Sigil ->
                Type -> Type
  substType' munif mrp mrow ms = go
   where
    go (PrimType pt) = PrimType pt
    go (Record rp r s) =
      let rp' = substRecPar' mrp rp
          r' = substRow' munif mrp mrow ms r
          s' = substSigil' ms s
       in Record rp' r' s'
    go (AbsType n s ts) =
      let s' = substSigil' ms s
          ts' = map go ts
       in  AbsType n s' ts'
    go (Variant r) = Variant $ substRow' munif mrp mrow ms r
    go (TypeVar v)      = TypeVar v
    go (TypeVarBang v)  = TypeVarBang v
    go (RecPar p c)     = RecPar p c
    go (RecParBang p c) = RecParBang p c
    go (Function t1 t2) = Function (go t1) (go t2)
    go t@(UnifVar v)    = fromMaybe t $ munif M.!? v
    go (Bang t)         = Bang $ go t
  
  substRecPar' :: M.Map VarName RecPar -> RecPar -> RecPar
  substRecPar' mrp rp@(UnknownParameter v) = fromMaybe rp $ mrp M.!? v
  substRecPar' _ rp = rp
  
  substSigil' :: M.Map VarName Sigil -> Sigil -> Sigil
  substSigil' ms s@(UnknownSigil v) = fromMaybe s $ ms M.!? v
  substSigil' _ s = s
  
  substRow' :: M.Map VarName Type ->
               M.Map VarName RecPar ->
               M.Map VarName Row ->
               M.Map VarName Sigil ->
               Row -> Row
  substRow' munif mrp mrow ms (Row m0 rvar) =
    let m1 = M.map (\(Entry n t tk) -> Entry n (substType' munif mrp mrow ms t) tk) m0
     in case rvar of
          (Just v) | Just (Row m' v') <- mrow M.!? v
                   -> Row (M.union m' m1) v'
          rvar     -> Row m1 rvar
  