content: {-# LANGUAGE DataKinds #-}
  {-# LANGUAGE GADTs #-}
  {-# LANGUAGE StandaloneDeriving #-}
  {-# LANGUAGE FlexibleInstances #-}
  {-# LANGUAGE MultiParamTypeClasses #-}
  {-# LANGUAGE DeriveGeneric #-}
  
  -- | This module provides the representation of theorems (predicates)
  -- | inside the type checker.
  module Language.Granule.Checker.Predicates where
  
  
  import Control.Monad.Trans.State.Strict
  import Data.List (intercalate, (\\))
  import GHC.Generics (Generic)
  
  import Language.Granule.Context
  
  import Language.Granule.Syntax.Helpers
  import Language.Granule.Syntax.Identifiers
  import Language.Granule.Syntax.FirstParameter
  import Language.Granule.Syntax.Pretty
  import Language.Granule.Syntax.Span
  import Language.Granule.Syntax.Type
  
  import Language.Granule.Utils
  
  import Data.Char (toLower)
  
  data Quantifier =
      -- | Universally quantification, e.g. polymorphic
      ForallQ
  
      -- | Unification variables
      | InstanceQ
  
      -- | Pi types
      | BoundQ
    deriving (Show, Eq)
  
  instance Pretty Quantifier where
    pretty ForallQ   = "∀"
    pretty InstanceQ = "[u]"
    pretty BoundQ    = "Π"
  
  -- Represent constraints generated by the type checking algorithm
  data Constraint =
    -- Coeffect resource algebra constraints
      Eq  Span Type Type Type
    | Neq Span Type Type Type
    | ApproximatedBy Span Type Type Type
    -- last argument is the kind of the coeffect
  
    -- (Least) upper bound; the last argument controls whether
    -- we do a check for leastness of the third argument (True) or not (False)
    -- Lub _ _ c1 c2 c3 means check that c1 <= c3 and c2 <= c3
    | Lub Span Coeffect Type Type Type Bool
  
    -- Used for user-predicates, and also effect types
    | Lt Span   Type Type -- Must be Nat kinded
    | Gt Span   Type Type -- Must be Nat kinded
    | LtEq Span Type Type -- Must be Nat kinded
    | GtEq Span Type Type -- Must be Nat kinded
  
    | Hsup Span Type Type Type
  
    deriving (Show, Eq, Generic)
  
  isEq :: Constraint -> Bool
  isEq (Eq _ _ _ _) = True
  isEq _ = False
  
  instance FirstParameter Constraint Span
  
  normaliseConstraint :: Constraint -> Constraint
  normaliseConstraint (Eq s c1 c2 t)   = Eq s (normalise c1) (normalise c2) t
  normaliseConstraint (Neq s c1 c2 t)  = Neq s (normalise c1) (normalise c2) t
  normaliseConstraint (Lub s c1 c2 c3 t b) = Lub s (normalise c1) (normalise c2) (normalise c3) t b
  normaliseConstraint (ApproximatedBy s c1 c2 t) = ApproximatedBy s (normalise c1) (normalise c2) t
  normaliseConstraint (Lt s c1 c2) = Lt s (normalise c1) (normalise c2)
  normaliseConstraint (Gt s c1 c2) = Gt s (normalise c1) (normalise c2)
  normaliseConstraint (LtEq s c1 c2) = LtEq s (normalise c1) (normalise c2)
  normaliseConstraint (GtEq s c1 c2) = GtEq s (normalise c1) (normalise c2)
  normaliseConstraint (Hsup s c1 c2 t) = Hsup s (normalise c1) (normalise c2) t
  
  instance Monad m => Freshenable m Constraint where
    freshen (Eq s' c1 c2 k) = do
      c1 <- freshen c1
      c2 <- freshen c2
      return $ Eq s' c1 c2 k
  
    freshen (Neq s' c1 c2 k) = do
      c1 <- freshen c1
      c2 <- freshen c2
      return $ Neq s' c1 c2 k
  
    freshen (ApproximatedBy s' c1 c2 t) = do
      c1 <- freshen c1
      c2 <- freshen c2
      return $ ApproximatedBy s' c1 c2 t
  
    freshen (Lub s' c1 c2 c3 t checkLeast) = do
      c1 <- freshen c1
      c2 <- freshen c2
      c3 <- freshen c3
      return $ Lub s' c1 c2 c3 t checkLeast
  
    freshen (Lt s c1 c2) = do
      c1 <- freshen c1
      c2 <- freshen c2
      return $ Lt s c1 c2
  
    freshen (Gt s c1 c2) = do
      c1 <- freshen c1
      c2 <- freshen c2
      return $ Gt s c1 c2
  
    freshen (LtEq s c1 c2) = LtEq s <$> freshen c1 <*> freshen c2
    freshen (GtEq s c1 c2) = GtEq s <$> freshen c1 <*> freshen c2
  
    freshen (Hsup s c1 c2 t) = do
      c1 <- freshen c1
      c2 <- freshen c2
      return $ Hsup s c1 c2 t
  
  -- Used to negate constraints
  newtype Neg a = Neg a
    deriving (Eq, Show)
  
  -- Used for pretty printing error messages with predicates in them
  prettyNegPred :: (?globals :: Globals) => Id -> Pred -> String
  -- Shorter message
  prettyNegPred defId (Con c) =
      "  When checking `" <> pretty defId <> "`, " <> message
    where
      message = toLower (head msg) : (tail msg)
      msg = pretty (Neg c)
  -- Long-winded message
  prettyNegPred defId p =
      "  The following constraint associated with `"
    <> pretty defId <> "` is false:\n\t  "
    <> pretty p
  
  instance Pretty (Neg Constraint) where
      pretty (Neg (Neq _ c1 c2 _)) =
        "Trying to prove that " <> pretty c1 <> " == " <> pretty c2
  
      pretty (Neg (Eq _ c1 c2 _)) =
        "Actual grade or index `" <> pretty c1 <>
        "` is not equal to specified grade or index `" <> pretty c2 <> "`"
  
      pretty (Neg (ApproximatedBy _ c1 c2 (TyCon k))) | internalName k == "Level" =
        pretty c2 <> " value cannot be moved to level " <> pretty c1
  
      pretty (Neg (ApproximatedBy _ c1 c2 k)) =
        case k of
          TyCon (internalName -> "Nat") ->
            "Expected " <> pretty c2 <> " uses, but instead there are " <> pretty c1 <> " actual uses."
          _ -> pretty c1 <> " is not approximatable by " <> pretty c2 <> " for type " <> pretty k
  
      pretty (Neg p@Lub{}) =
        "Trying to prove negation of statement: " ++ pretty p
  
      pretty (Neg (Lt _ c1 c2)) =
        "Trying to prove false statement: (" <> pretty c1 <> " < " <> pretty c2 <> ")"
  
      pretty (Neg (Gt _ c1 c2)) =
        "Trying to prove false statement: (" <> pretty c1 <> " > " <> pretty c2 <> ")"
  
      pretty (Neg (LtEq _ c1 c2)) =
        "Trying to prove false statement: (" <> pretty c1 <> " ≤ " <> pretty c2 <> ")"
  
      pretty (Neg (GtEq _ c1 c2)) =
        "Trying to prove false statement: (" <> pretty c1 <> " ≥ " <> pretty c2 <> ")"
  
      pretty (Neg (Hsup _ c1 c2 t)) =
        "Cannot split resource requirements into " <> pretty c1 <> " and " <> pretty c2 <>
        " in a pattern match. \nTrying to prove false statement: (" <> pretty c1 <> " ⨱ " <> pretty c2 <> ")"
  
  instance Pretty [Constraint] where
      pretty constr =
        "---\n" <> (intercalate "\n" . map pretty $ constr)
  
  instance Pretty Constraint where
      pretty (Eq _ c1 c2 _) =
        "(" <> pretty c1 <> " = " <> pretty c2 <> ")" -- @" <> show s
  
      pretty (Neq _ c1 c2 _) =
          "(" <> pretty c1 <> " ≠ " <> pretty c2 <> ")" -- @" <> show s
  
      pretty (ApproximatedBy _ c1 c2 k) =
        case k of
          -- Nat is discrete
          TyCon (internalName -> "Nat") ->
              "(" <> pretty c1 <> " = " <> pretty c2 <> ")"
          -- Sets (subset)
          TyApp (TyCon (internalName -> "Set")) _ ->
              "(" <> pretty c1 <> " ⊆ " <> pretty c2 <> ")"
          -- Sets (superset)
          TyApp (TyCon (internalName -> "SetOp")) _ ->
              "(" <> pretty c1 <> " ⊇ " <> pretty c2 <> ")"
  
          _ -> "(" <> pretty c1 <> " ≤ " <> pretty c2 <> ")" -- <> " @ " <> pretty k
  
      pretty (Lub _ c1 c2 c3 _ _) =
        "(" <> pretty c1 <> " ⊔ " <> pretty c2 <> " = " <> pretty c3 <> ")"
  
      pretty (Lt _ c1 c2) =
        "(" <> pretty c1 <> " < " <> pretty c2 <> ")"
  
      pretty (Gt _ c1 c2) =
        "(" <> pretty c1 <> " > " <> pretty c2 <> ")"
  
      pretty (LtEq _ c1 c2) =
        "(" <> pretty c1 <> " ≤ " <> pretty c2 <> ")"
  
      pretty (GtEq _ c1 c2) =
        "(" <> pretty c1 <> " ≥ " <> pretty c2 <> ")"
  
      pretty (Hsup _ c1 c2 t) =
        "(" <> pretty c1 <> " ⨱ " <> pretty c2 <> ")"
  
  varsConstraint :: Constraint -> [Id]
  varsConstraint (Eq _ c1 c2 _) = freeVars c1 <> freeVars c2
  varsConstraint (Neq _ c1 c2 _) = freeVars c1 <> freeVars c2
  varsConstraint (Lub _ c1 c2 c3 _ _) = freeVars c1 <> freeVars c2 <> freeVars c3
  varsConstraint (ApproximatedBy _ c1 c2 _) = freeVars c1 <> freeVars c2
  varsConstraint (Lt _ c1 c2) = freeVars c1 <> freeVars c2
  varsConstraint (Gt _ c1 c2) = freeVars c1 <> freeVars c2
  varsConstraint (LtEq _ c1 c2) = freeVars c1 <> freeVars c2
  varsConstraint (GtEq _ c1 c2) = freeVars c1 <> freeVars c2
  varsConstraint (Hsup _ c1 c2 t) = freeVars c1 <> freeVars c2
  
  -- Represents a predicate generated by the type checking algorithm
  data Pred where
      Conj :: [Pred] -> Pred
      Disj :: [Pred] -> Pred
      Impl :: Ctxt Kind -> Pred -> Pred -> Pred
      Con  :: Constraint -> Pred
      NegPred  :: Pred -> Pred
      Exists :: Id -> Kind -> Pred -> Pred
  
  instance Term Pred where
    freeVars (Conj ps) = concatMap freeVars ps
    freeVars (Disj ps) = concatMap freeVars ps
    freeVars (Impl bounds p1 p2) = (freeVars p1 <> freeVars p2) \\ map fst bounds
    freeVars (Con c) = varsConstraint c
    freeVars (NegPred p) = freeVars p
    freeVars (Exists x _ p) = freeVars p \\ [x]
  
  boundVars :: Pred -> [Id]
  boundVars (Conj ps) = concatMap boundVars ps
  boundVars (Disj ps) = concatMap boundVars ps
  boundVars (Impl bounds p1 p2) = map fst bounds ++ (boundVars p1 ++ boundVars p2)
  boundVars (NegPred p) = boundVars p
  boundVars (Exists x _ p) = x : boundVars p
  boundVars (Con _) = []
  
  instance (MonadFail m) => Freshenable m Pred where
    freshen (Conj ps) = do
      ps' <- mapM freshen ps
      return $ Conj ps'
  
    freshen (Disj ps) = do
      ps' <- mapM freshen ps
      return $ Disj ps'
  
    freshen (NegPred p) = do
      p' <- freshen p
      return $ NegPred p'
  
    freshen (Exists v k p) = do
      st <- get
  
      -- Create a new binding name for v
      let v' = internalName v <> "-e" <> show (counter st)
  
      -- Updated freshener state
      put (st { tyMap = (internalName v, v') : tyMap st
            , counter = counter st + 1 })
  
      -- Freshen the rest of the predicate
      p' <- freshen p
      -- Freshening now out of scope
      removeFreshenings [Id (internalName v) v']
  
      return $ Exists (Id (internalName v) v') k p'
  
    freshen (Impl [] p1 p2) = do
      p1' <- freshen p1
      p2' <- freshen p2
      return $ Impl [] p1' p2'
  
    freshen (Impl ((v, kind):vs) p p') = do
      st <- get
  
      -- Freshen the variable bound here
      let v' = internalName v <> "-" <> show (counter st)
      put (st { tyMap = (internalName v, v') : tyMap st
              , counter = counter st + 1 })
  
      -- Freshen the rest
      (Impl vs' pf pf') <- freshen (Impl vs p p')
      -- Freshening now out of scope
      removeFreshenings [Id (internalName v) v']
  
      return $ Impl ((Id (internalName v) v', kind):vs') pf pf'
  
    freshen (Con cons) = do
      cons' <- freshen cons
      return $ Con cons'
  
  deriving instance Show Pred
  deriving instance Eq Pred
  
  -- Fold operation on a predicate
  predFold ::
       ([a] -> a)
    -> ([a] -> a)
    -> (Ctxt Kind -> a -> a -> a)
    -> (Constraint -> a)
    -> (a -> a)
    -> (Id -> Kind -> a -> a)
    -> Pred
    -> a
  predFold c d i a n e (Conj ps)   = c (map (predFold c d i a n e) ps)
  predFold c d i a n e (Disj ps)   = d (map (predFold c d i a n e) ps)
  predFold c d i a n e (Impl ctxt p p') = i ctxt (predFold c d i a n e p) (predFold c d i a n e p')
  predFold _ _ _ a _  _ (Con cons)  = a cons
  predFold c d i a n e (NegPred p) = n (predFold c d i a n e p)
  predFold c d i a n e (Exists x t p) = e x t (predFold c d i a n e p)
  
  -- Fold operation on a predicate (monadic)
  predFoldM :: Monad m =>
       ([a] -> m a)
    -> ([a] -> m a)
    -> (Ctxt Kind -> a -> a -> m a)
    -> (Constraint -> m a)
    -> (a -> m a)
    -> (Id -> Kind -> a -> m a)
    -> Pred
    -> m a
  predFoldM c d i a n e (Conj ps)   = do
    ps <- mapM (predFoldM c d i a n e) ps
    c ps
  
  predFoldM c d i a n e (Disj ps)   = do
    ps <- mapM (predFoldM c d i a n e) ps
    d ps
  
  predFoldM c d i a n e (Impl localVars p p') = do
    p <- predFoldM c d i a n e p
    p' <- predFoldM c d i a n e p'
    i localVars p p'
  
  predFoldM _ _ _ a _ _ (Con cons)  =
    a cons
  
  predFoldM c d i a n e (NegPred p) =
    predFoldM c d i a n e p >>= n
  
  predFoldM c d i a n e (Exists x t p) =
    predFoldM c d i a n e p >>= e x t
  
  instance Pretty [Pred] where
    pretty ps =
      "Size = " <> show (length ps) <> "\n" <>
      (intercalate "\n" (map (\p -> " - " <> pretty p) ps))
  
  instance Pretty Pred where
    pretty =
      (predFold
       (\x -> if null x then "T" else intercalate " ∧ " x)
       (intercalate " ∨ ")
       (\ctxt p q ->
           (if null ctxt then "" else "∀ {" <> pretty' ctxt <> "} . ")
        <> "((" <> p <> ") -> " <> q <> ")")
        pretty
        (\p -> "¬(" <> p <> ")")
        (\x t p -> "∃ " <> pretty x <> " : " <> pretty t <> " . " <> p))
      . preFilterImplCtxts
      where
        preFilterImplCtxts =
          predFold Conj Disj
            (\ctxt p q -> Impl (filter (\(id, k) -> ((id `elem` freeVars p) || (id `elem` freeVars q))) ctxt) p q)
            Con NegPred Exists
  
        pretty' xs = intercalate ", " (map prettyBinding xs)
        prettyBinding (id, k) = pretty id <> " : " <> pretty k
  
  -- | Whether the predicate is empty, i.e. contains no constraints
  isTrivial :: Pred -> Bool
  isTrivial = predFold and or (\_ lhs rhs -> rhs) (const False) id (\_ _ p -> p)
  
  -- TODO: consider replacing this with use of `substitute`?
  -- given a context mapping coeffect type variables to coeffect typ,
  -- then rewrite a set of constraints so that any occruences of the kind variable
  -- are replaced with the coeffect type
  rewriteBindersInPredicate :: Ctxt (Type, Quantifier) -> Pred -> Pred
  rewriteBindersInPredicate ctxt =
      predFold
        Conj
        Disj
        Impl
        (\c -> Con $ foldr (uncurry updateConstraint) c ctxt)
        NegPred
        existsCase
    where
      existsCase :: Id -> Kind -> Pred -> Pred
      existsCase var (TyVar kvar) p =
        Exists var k' p
          where
            k' = case lookup kvar ctxt of
                    Just (ty, _) -> ty
                    Nothing -> TyVar kvar
      existsCase var k p = Exists var k p
  
      -- `updateConstraint v k c` rewrites any occurence of the kind variable
      -- `v` in the constraint `c` with the kind `k`
      updateConstraint :: Id -> (Type, Quantifier) -> Constraint -> Constraint
      updateConstraint ckindVar (ckind, _) (Eq s c1 c2 k) =
        Eq s (updateCoeffect ckindVar ckind c1) (updateCoeffect ckindVar ckind c2)
          (case k of
            TyVar ckindVar' | ckindVar == ckindVar' -> ckind
            _ -> k)
      updateConstraint ckindVar (ckind, _) (Neq s c1 c2 k) =
              Neq s (updateCoeffect ckindVar ckind c1) (updateCoeffect ckindVar ckind c2)
                (case k of
                  TyVar ckindVar' | ckindVar == ckindVar' -> ckind
                  _ -> k)
  
      updateConstraint ckindVar (ckind, _) (ApproximatedBy s c1 c2 k) =
        ApproximatedBy s (updateCoeffect ckindVar ckind c1) (updateCoeffect ckindVar ckind c2)
          (case k of
            TyVar ckindVar' | ckindVar == ckindVar' -> ckind
            _  -> k)
  
      updateConstraint ckindVar (ckind, _) (Lub s c1 c2 c3 k checkLeast) =
        Lub s (updateCoeffect ckindVar ckind c1) (updateCoeffect ckindVar ckind c2) (updateCoeffect ckindVar ckind c3)
          (case k of
            TyVar ckindVar' | ckindVar == ckindVar' -> ckind
            _  -> k) checkLeast
  
      updateConstraint ckindVar (ckind, _) (Lt s c1 c2) =
          Lt s (updateCoeffect ckindVar ckind c1) (updateCoeffect ckindVar ckind c2)
  
      updateConstraint ckindVar (ckind, _) (Gt s c1 c2) =
          Gt s (updateCoeffect ckindVar ckind c1) (updateCoeffect ckindVar ckind c2)
  
      updateConstraint ckindVar (ckind, _) (GtEq s c1 c2) =
          GtEq s (updateCoeffect ckindVar ckind c1) (updateCoeffect ckindVar ckind c2)
  
      updateConstraint ckindVar (ckind, _) (LtEq s c1 c2) =
          LtEq s (updateCoeffect ckindVar ckind c1) (updateCoeffect ckindVar ckind c2)
  
      updateConstraint ckindVar (ckind, _) (Hsup s c1 c2 t) =
          Hsup s (updateCoeffect ckindVar ckind c1) (updateCoeffect ckindVar ckind c2) t
  
      -- `updateCoeffect v k c` rewrites any occurence of the kind variable
      -- `v` with the kind `k` in the coeffect `c`
      updateCoeffect :: Id -> Type -> Type -> Type
      updateCoeffect ckindVar ckind (TyGrade (Just (TyVar ckindVar')) c) | ckindVar == ckindVar' = TyGrade (Just ckind) c
      updateCoeffect ckindVar ckind (TySig c (TyVar ckindVar')) | ckindVar == ckindVar' = TySig c ckind
      updateCoeffect ckindVar ckind (TyInfix op c1 c2) =
        TyInfix op (updateCoeffect ckindVar ckind c1) (updateCoeffect ckindVar ckind c2)
      updateCoeffect _ _ c = c
  
  ------------------------------------------------------------------------------------------------------------
  -- Zipper formulation of predicates (more expressive for working in a context where we need to solve on partial formulae)
  
  data PredContext =
      Top
    | ExistsHere     Id Kind PredContext
    | NegPredHere    PredContext
    | ConjHere       [Pred] PredContext
    | DisjHere       [Pred] PredContext
    | ImplAntecedent (Ctxt Kind) PredContext
    | ImplConsequent (Ctxt Kind) Pred PredContext
    deriving (Show, Eq) -- for debugging
  
  
  -----    ImplConsequent [("x", Nat)] p path
  -----
  ----     represents
  -----
  -----       .
  -----       .  path
  -----       .
  ----     forall x : Nat . _ -> _
  ----     /   \
  ---     p    [ ]
  
  ---    ImplAntecedent [("x", Nat)] path
  
  ---      .
  ---      .  path
  ---      .
  ---   forall x : Nat . ->
  ---      /      \
  ---    [ ]      ?
  
  
  ----   fromPredicateContext (ImplConsequent [("x", Nat)] p (ConjHere [Con c] Top))
  
  ---    -> Conj [Impl [("x", Nat)] p (Conj []), Con c]
  
  
  fromPredicateContext :: PredContext -> Pred
  -- fill the 'hole' with True (which is just an empty conjunction)
  fromPredicateContext path = fromPredicateZipper path (Conj [])
    where
      fromPredicateZipper :: PredContext -> Pred -> Pred
      fromPredicateZipper Top                           p = p
      fromPredicateZipper (ExistsHere v k path)         p = fromPredicateZipper path (Exists v k p)
      fromPredicateZipper (NegPredHere path)            p = fromPredicateZipper path (NegPred p)
      fromPredicateZipper (ConjHere ps path)            p = fromPredicateZipper path (Conj (p : ps))
      fromPredicateZipper (DisjHere ps path)            p = fromPredicateZipper path (Disj (p : ps))
      fromPredicateZipper (ImplAntecedent ctxt path)    p = fromPredicateZipper path (Impl ctxt p (Conj []))
      fromPredicateZipper (ImplConsequent ctxt pa path) p = fromPredicateZipper path (Impl ctxt pa p)
  
  newImplication :: Ctxt Kind -> PredContext -> PredContext
  newImplication ctxt path = ImplAntecedent ctxt path
  
  moveToConsequent :: PredContext -> Either String PredContext
  moveToConsequent path = rollup path (Conj [])
    where
      rollup (ExistsHere v k path)         p = rollup path (Exists v k p)
      rollup (NegPredHere path)            p = rollup path (NegPred p)
      rollup (ConjHere ps path)            p = rollup path (Conj (p:ps))
      rollup (DisjHere ps path)            p = rollup path (Disj (p:ps))
      rollup (ImplConsequent ctxt pa path) p = rollup path (Impl ctxt pa p)
      rollup (ImplAntecedent ctxt path)    p = Right $ ImplConsequent ctxt p path
      rollup Top                           p =
        Left $ "Bug: reached top of formula without hitting an implication, rolled up " ++ show p
  
  moveToNewConjunct :: PredContext -> PredContext
  moveToNewConjunct path = rollup path (Conj [])
    where
      rollup (ExistsHere v k path)         p = rollup path (Exists v k p)
      rollup (NegPredHere path)            p = rollup path (NegPred p)
      rollup (ConjHere ps path)            p = ConjHere (p:ps) path
      rollup (DisjHere ps path)            p = rollup path (Disj (p:ps))
      rollup (ImplConsequent ctxt pa path) p = rollup path (Impl ctxt pa p)
      rollup (ImplAntecedent ctxt path)    p = rollup path (Impl ctxt p (Conj []))
      rollup Top                           p = ConjHere [] Top
  
  addConstraintViaConjunction :: Constraint -> PredContext -> PredContext
  addConstraintViaConjunction c (ConjHere ps path) =
    ConjHere (addToConjunct (Con c) ps) path
  addConstraintViaConjunction c path = ConjHere [Con c] path
  
  addPredicateViaConjunction :: Pred -> PredContext -> PredContext
  addPredicateViaConjunction p (ConjHere ps path) =
    ConjHere (addToConjunct p ps) path
  addPredicateViaConjunction p path               = ConjHere [p] path
  
  addToConjunct :: Pred -> [Pred] -> [Pred]
  addToConjunct p ((Conj ps):pss) =
    Conj (p:ps) : pss
  addToConjunct p ps = p:ps
  