content: -- |
  -- Module      : Minigent.TC.ConstraintGen
  -- Copyright   : (c) Data61 2018-2019
  --                   Commonwealth Science and Research Organisation (CSIRO)
  --                   ABN 41 687 119 230
  -- License     : BSD3
  --
  -- The constraint generator portion of the type inference.
  --
  -- May be used qualified or unqualified.
  {-# LANGUAGE GeneralizedNewtypeDeriving #-}
  module Minigent.TC.ConstraintGen (CG, runCG, cg, cgFunction) where
  
  import Minigent.Syntax
  import Minigent.Syntax.Utils
  import qualified Minigent.Syntax.Utils.Row as Row
  import Minigent.Fresh
  import Minigent.Environment
  
  import Control.Monad.Reader
  import Control.Monad.Fail
  import Control.Monad.State.Strict
  
  import qualified Data.Map as M
  import Data.Stream (Stream)
  
  import Minigent.Syntax.PrettyPrint
  import Debug.Trace
  
  -- | A monad that is a combination of a state monad for the current type context,
  --   a reader monad for the global environments, and fresh variables.
  newtype CG a = CG (StateT (Context Type) (ReaderT GlobalEnvironments (FreshT VarName IO)) a)
          deriving ( Monad, Applicative, Functor, MonadState (Context Type)
                   , MonadFresh VarName, MonadReader GlobalEnvironments
                   )
  
  -- | Given an initial context and global environments, run a 'CG' computation.
  runCG :: CG a -> Context Type -> GlobalEnvironments -> FreshT VarName IO a
  runCG (CG x) ctx glb = runReaderT (evalStateT x ctx) glb
  
  -- | The constraint generation relation for expressions as given in the language
  --   definition.
  --
  --   @cg e tau@ returns @(e', c)@, where @e'@ is an annotated version of @e@ that contains
  --   unknowns (unification variables) that, when substituted by a satisfying assignment to @c@,
  --   produces a well-typed epxression.
  cg :: Expr -> Type -> CG (Expr, Constraint)
  cg e tau = case e of
  
    (Var v) -> do
      (rho, c) <- lookupVar v
      withSig (Var v, c :&: rho :< tau)
  
    (Sig e1 tau') -> do
      (e1', c) <- cg e1 tau'
      withSig (Sig e1' tau', c :&: tau' :< tau)
  
    (Apply e1 e2) -> do
      alpha <- UnifVar <$> fresh
      (e1', c1) <- cg e1 (Function alpha tau)
      (e2', c2) <- cg e2 alpha
      withSig (Apply e1' e2', c1 :&: c2)
  
    (Let x e1 e2) -> do
      alpha <- UnifVar <$> fresh
      (e1', c1) <- cg e1 alpha
      modify (push (x, alpha))
      (e2', c2) <- cg e2 tau
      used <- topUsed <$> get
      let c3 = if not used then Drop alpha else Sat
      modify pop
      withSig (Let x e1' e2', c1 :&: c2 :&: c3)
  
    (Literal l) -> do
      let c = case l of
                BoolV _ -> PrimType Bool :< tau
                UnitV   -> PrimType Unit :< tau
                IntV  l -> l :<=: tau
      withSig (e, c)
  
    (If e1 e2 e3) -> do
      (e1', c1) <- cg e1 (PrimType Bool)
      g2 <- get
      (e2', c2) <- cg e2 tau
      g3 <- get
      put g2
      (e3', c3) <- cg e3 tau
      g3' <- get
      let (as, g4) = reconcile g3 g3'
          c4       = conjunction (map Drop as)
      put g4
      withSig (If e1' e2' e3', c1 :&: c2 :&: c3 :&: c4)
  
    (PrimOp Not [e]) -> do
      (e', c) <- cg e (PrimType Bool)
      withSig (PrimOp Not [e'], c :&: PrimType Bool :< tau)
  
    (PrimOp o [e1,e2])
      | o `elem` boolOps -> do
        (e1', c1) <- cg e1 (PrimType Bool)
        (e2', c2) <- cg e2 (PrimType Bool)
        withSig (PrimOp o [e1', e2'], c1 :&: c2 :&: PrimType Bool :< tau)
      | o `elem` compOps -> do
          alpha <- UnifVar <$> fresh
          (e1', c1) <- cg e1 alpha
          (e2', c2) <- cg e2 alpha
          withSig (PrimOp o [e1', e2'], 0 :<=: alpha :&: c1 :&: c2 :&: PrimType Bool :< tau)
      | otherwise -> do
          (e1', c1) <- cg e1 tau
          (e2', c2) <- cg e2 tau
          withSig (PrimOp o [e1', e2'], 0 :<=: tau :&: c1 :&: c2)
  
    (TypeApp f ts) -> do
      pt <- M.lookup f . types <$> ask
      let (vs, cs, t) = case pt of
                          Just (Forall vs cs t) -> (vs, cs, t)
                          _ -> error "cg: TypeApp did not have a type in types"
      as <- freshes (length vs - length ts)
      let ts'   = ts ++ map UnifVar as
          subst = zip vs ts'
          cs'   = map (constraintTypes (traverseType (substTVs subst))) cs
      withSig (TypeApp f ts', traverseType (substTVs subst) t :< tau :&: conjunction cs')
  
    (Con n e) -> do
      alpha <- UnifVar <$> fresh
      row <- Row.incomplete [Entry n alpha False]
      (e', c) <- cg e alpha
      withSig (Con n e', Variant row :< tau :&: c)
  
    (Case e1 k x e2 y e3) -> do
      beta <- UnifVar <$> fresh
      row <- Row.incomplete [Entry k beta False]
      (e1', c1) <- cg e1 (Variant row)
      g2 <- get
      modify (push (x, beta))
      (e2', c2) <- cg e2 tau
      xUsed <- topUsed <$> get
      let c6 = if xUsed then Sat else Drop beta
      modify pop
      g3 <- get
      put g2
      let row' = Row.take k row
      modify (push (y, Variant row'))
      (e3', c3) <- cg e3 tau
      yUsed <- topUsed <$> get
      let c7 = if yUsed then Sat else Drop (Variant row')
      modify pop
      g3' <- get
      let (as, g4) = reconcile g3 g3'
          c4       = conjunction (map Drop as)
      put g4
      withSig (Case e1' k x e2' y e3', c1 :&: c2 :&: c3 :&: c4 :&: c6 :&: c7)
  
  
    (Esac e1 k x e2) -> do
      beta <- UnifVar <$> fresh
      row <- Row.incomplete [Entry k beta False]
      (e1', c1) <- cg e1 (Variant row)
      modify (push (x, beta))
      (e2', c2) <- cg e2 tau
      xUsed <- topUsed <$> get
      let c4 = if xUsed then Sat else Drop beta
          c3 = Exhausted (Variant (Row.take k row))
      modify pop
      withSig (Esac e1' k x e2', c1 :&: c2 :&: c3 :&: c4)
  
    (LetBang ys x e1 e2) -> do
      alpha <- UnifVar . (++"letba") <$> fresh
      modify (alter ys (\(t,u) -> (Bang t, 0)))
      (e1', c1) <- cg e1 alpha
      bangRhos' <- factor ys <$> get
      let c3 = conjunction (map Drop (unused bangRhos'))
      modify (alter ys (\(Bang t, u) -> (t,0)))
      modify (push (x, alpha))
      (e2', c2) <- cg e2 tau
      xUsed <- topUsed <$> get
      let c5 = if xUsed then Sat else Drop alpha
      modify pop
      rhos' <- factor ys <$> get
      let c4 = conjunction (map Drop (unused rhos'))
          c6 = Escape alpha
      withSig (LetBang ys x e1' e2', c1 :&: c2 :&: c3 :&: c4 :&: c5 :&: c6)
  
    (Member e f) -> do
      row <- Row.incomplete [Entry f tau False]
      sigil <- fresh
      recPar <- UnknownParameter <$> fresh
      let alpha = Record recPar row (UnknownSigil sigil)
      (e', c1) <- cg e alpha
      let c2 = Drop (Record recPar (Row.take f row) (UnknownSigil sigil))
      withSig (Member e' f, c1 :&: c2)
  
    -- Remaining record, field name, extracted contents var, record extrating from, following expression 
    (Take x f y e1 e2) -> do
      beta <- UnifVar <$> fresh
      row <- Row.incomplete [Entry f beta False]
      sigil  <- fresh
      recPar <- UnknownParameter <$> fresh
      let alpha = Record recPar row (UnknownSigil sigil)
      let c0    = UnboxedNoRecurse recPar (UnknownSigil sigil)
  
      (e1', c1) <- cg e1 alpha
      modify (push (y, beta))
      modify (push (x, Record recPar (Row.take f row) (UnknownSigil sigil)))
      (e2', c2) <- cg e2 tau
      xUsed <- topUsed <$> get
      let c3 = if xUsed then Sat else Drop (Record recPar (Row.take f row) (UnknownSigil sigil))
      modify pop
      yUsed <- topUsed <$> get
      let c4 = if yUsed then Sat else Drop beta
      modify pop
      withSig (Take x f y e1' e2', c0 :&: c1 :&: c2 :&: c3 :&: c4)-- :&: c5)
  
    (Put e1 f e2) -> do
      beta <- UnifVar <$> fresh
      sigil <- fresh
      recPar <- UnknownParameter <$> fresh
      row  <- Row.incomplete [Entry f beta True]
  
      let alpha = Record recPar row (UnknownSigil sigil)
      (e1', c1) <- cg e1 alpha
      (e2', c2) <- cg e2 beta
  
      let c0 = UnboxedNoRecurse recPar (UnknownSigil sigil)
  
      let c3 = Record recPar (Row.put f row) (UnknownSigil sigil) :< tau
  
      withSig (Put e1' f e2', c0 :&: c1 :&: c2 :&: c3)
  
    (Struct fs) -> do
      (fs', ts, cs) <- cgStruct fs
      withSig (Struct fs', conjunction cs :&: Record None (Row.fromList ts) Unboxed :< tau )
  
    where
  
      cgStruct :: [(FieldName, Expr)] -> CG ([(FieldName, Expr)], [Entry], [Constraint])
      cgStruct [] = return ([], [], [])
      cgStruct ((f,e):fs) = do
        (fs', ts', cs') <- cgStruct fs
        alpha <- UnifVar <$> fresh
        (e', c) <- cg e alpha
        return ((f,e'):fs', (Entry f alpha False):ts', c:cs')
  
  
      withSig :: (Expr, Constraint) -> CG (Expr, Constraint)
      withSig (e, c) = pure (Sig e tau, c)
  
      -- | Looks up a variable and increases it's usage count, adding the
      --   constraint that it is shareable if it's been used more than once
      lookupVar :: VarName -> CG (Type, Constraint)
      lookupVar v = do
        (rho, used, ctx') <- use v <$> get
        put ctx'
        return (rho, if used then Share rho else Sat)
  
  
  -- | Used for constraint generation for top-level functions.
  --   Given a function name, argument name and a function body expression,
  --   return an annotated function body along with the constraint that would make
  --   it well typed. Also included in the first component of the return value
  --   are the axioms (constraints placed by the user in the type signature)
  --   about polymorphic type variables.
  cgFunction :: FunName -> VarName -> Expr -> CG ([Constraint], Expr, Constraint)
  cgFunction f x e = do
    alpha <- UnifVar <$> fresh
    beta  <- UnifVar <$> fresh
    let proposedType = Function alpha beta
    modify (push (x,alpha))
    (e', c) <- cg e beta
    used <- topUsed <$> get
    let c' = if used then Sat else Drop alpha
    modify pop
    envs <- ask
    let (c'',axs) = case M.lookup f (types envs) of
                         Nothing -> (Sat, [])
                         Just (Forall vs cs tau) -> (proposedType :< tau, cs)
    -- Inferred constraints for return type 
    -- && proposed function type is subtype of inferred function type
    -- && Argument to function is used or droppable
    pure (axs, e', c :&: c'' :&: c')
  