content: -- |
  -- Module      : Minigent.Syntax.Utils
  -- Copyright   : (c) Data61 2018-2019
  --                   Commonwealth Science and Research Organisation (CSIRO)
  --                   ABN 41 687 119 230
  -- License     : BSD3
  --
  -- This module provides various miscellaneous utility functions for querying
  -- and manipulating syntax.
  --
  -- It expects to be imported unqualified.
  module Minigent.Syntax.Utils
    ( -- * Operators
      operators
    , -- ** Operator categories
      -- | The various syntactic precendence categories of binary operators
      prodOps
    , sumOps
    , compOps
    , boolOps
    , -- * Constraints
      flattenConstraint
    , conjunction
    , constraintTypes
    , -- * Types
      -- ** Applying rewrites
      traverseType
    , normaliseType
    , sameRecursive
    , unroll
    , mapRecPars
    , -- ** Rewrites
      substUV
    , substRowV
    , substSigilV
    , substTV
    , substUVs
    , substTVs
    , substRecPar
    , -- ** Queries for type inference
      fits
    , unorderedType
    , typeUVs
    , typeVariables
    , recursiveParameterNames
    , rigid
    , rootUnifVar
    , -- * Entries
      entryTypes
    , -- * Sigils
      sigilsCompatible
    , -- * Expressions
      exprTypes
    , -- * Fresh Unification Variables
      unifVars
    , withUnifVars
    )
  where
  
  import           Minigent.Syntax
  import           Minigent.Fresh
  import qualified Minigent.Syntax.Utils.Rewrite as RW
  import qualified Minigent.Syntax.Utils.Row     as Row
  
  
  import           Control.Applicative
  import           Control.Monad                  ( guard )
  import           Data.Maybe                     ( fromMaybe
                                                  , maybeToList
                                                  , isNothing
                                                  )
  
  import qualified Data.Stream                   as S
  import qualified Data.Map                      as M
  
  -- | Returns true iff the given argument type is not subject to subtyping. That is, if @a :\< b@
  --   (subtyping) is equivalent to @a :=: b@ (equality), then this function returns true.
  --
  --   At least for now, this returns true for all types but variants, records and functions.
  unorderedType :: Type -> Bool
  unorderedType (Record{}  ) = False
  unorderedType (Variant{} ) = False
  unorderedType (Function{}) = False
  unorderedType t            = rigid t
  
  -- | Return all of the unification type variables inside a type.
  typeUVs :: Type -> [VarName]
  typeUVs (UnifVar v) = [v]
  typeUVs (Record n r s) = concatMap (\(Entry _ t _) -> typeUVs t) (Row.entries r)
                      ++ maybe [] pure (rowVar r)
                      ++ (case s of UnknownSigil     s' -> [s']; _ -> [])
                      ++ (case n of UnknownParameter n' -> [n']; _ -> [])
  typeUVs (Variant r)  = concatMap (\(Entry _ t _) -> typeUVs t) (Row.entries r)
                      ++ maybe [] pure (rowVar r)
  typeUVs (AbsType _ _ ts) = concatMap typeUVs ts
  typeUVs (Function t1 t2) = typeUVs t1 ++ typeUVs t2
  typeUVs (Bang t        ) = typeUVs t
  typeUVs _                = []
  
  -- | Return all of the (rigid, non-unification) type variables in a type. Does not include mu variables
  typeVariables :: Type -> [VarName]
  typeVariables t = typeVariables' t []
   where
      -- Ensures recursive parameters are not included in type variables
    typeVariables' :: Type -> [VarName] -> [VarName]
    typeVariables' (TypeVar     v) mvs = if elem v mvs then [] else [v]
    typeVariables' (TypeVarBang v) mvs = if elem v mvs then [] else [v]
    typeVariables' (Record mt r _) mvs = concatMap
      (\(Entry _ t _) -> typeVariables' t ((case mt of Rec x -> [x]; _ -> []) ++ mvs))
      (Row.entries r)
    typeVariables' (Variant r) mvs      = concatMap (\(Entry _ t _) -> typeVariables' t mvs) (Row.entries r)
    typeVariables' (AbsType _ _ ts) mvs = concatMap (\x -> typeVariables' x mvs) ts
    typeVariables' (Function t1 t2) mvs = typeVariables' t1 mvs ++ typeVariables' t2 mvs
    typeVariables' (Bang t        ) mvs = typeVariables' t mvs
    typeVariables' _                _   = []
  
  recursiveParameterNames :: Type -> [VarName]
  recursiveParameterNames (Record mt r _) = case mt of Rec x -> [x]; _ -> []
    ++ concatMap (\(Entry _ t _) -> recursiveParameterNames t) (Row.entries r)
  recursiveParameterNames (Variant r) =
    concatMap (\(Entry _ t _) -> recursiveParameterNames t) (Row.entries r)
  recursiveParameterNames (AbsType _ _ ts) = concatMap recursiveParameterNames ts
  recursiveParameterNames (Function t1 t2) = recursiveParameterNames t1 ++ recursiveParameterNames t2
  recursiveParameterNames (Bang t        ) = recursiveParameterNames t
  recursiveParameterNames _                = []
  
  
  -- | Returns @True@ unless the given type is a unification variable or a type operator
  --   applied to a unification variable.
  rigid :: Type -> Bool
  rigid (UnifVar _)  = False
  rigid (Bang _)     = False
  rigid (Record _ r _) = not $ Row.justVar r
  rigid (Variant r)  = not $ Row.justVar r
  rigid _            = True
  
  -- | Return the unification variable in a non-rigid type.
  --   If the type is rigid, then returns @Nothing@.
  rootUnifVar :: Type -> Maybe VarName
  rootUnifVar (UnifVar n   ) = Just n
  rootUnifVar (Bang    n   ) = rootUnifVar n
  rootUnifVar (Variant r   ) = rowVar r
  rootUnifVar (Record _ r s) = rowVar r
  rootUnifVar _              = Nothing
  
  -- | A table of all operators, mapping string representations
  --   to their 'Op' values.
  operators :: [(String, Op)]
  operators =
    [ ("+" , Plus)
    , ("*" , Times)
    , ("-" , Minus)
    , ("/" , Divide)
    , ("%" , Mod)
    , ("<" , Less)
    , (">" , Greater)
    , ("==", Equal)
    , ("/=", NotEqual)
    , ("<=", LessEqual)
    , (">=", GreaterEqual)
    , ("&&", And)
    , ("||", Or)
    , ("~" , Not)
    ]
  
  prodOps, sumOps, compOps, boolOps :: [Op]
  prodOps = [Times, Divide, Mod]
  sumOps = [Plus, Minus]
  compOps = [Equal, NotEqual, Greater, Less, GreaterEqual, LessEqual]
  boolOps = [And, Or, Not]
  
  -- | Given a constraint, flatten it out to remove all conjunctions,
  --   returning a list of conjunction-free constraints.
  flattenConstraint :: Constraint -> [Constraint]
  flattenConstraint (a :&: b) = flattenConstraint a ++ flattenConstraint b
  flattenConstraint x         = [x]
  
  -- | Given a list of constraints, combine them into one constraint
  --   using constraint conjunction.
  conjunction :: [Constraint] -> Constraint
  conjunction = foldr (:&:) Sat
  
  -- | Given a function that acts on 'Type' values, produce a function
  --   that acts on the type inside 'Entry' values.
  entryTypes :: (Type -> Type) -> Entry -> Entry
  entryTypes func (Entry f t k) = Entry f (func t) k
  
  -- | Given a function that acts on 'Type' values, produce a function
  --   that acts on the types inside 'Constraint' values.
  constraintTypes :: (Type -> Type) -> Constraint -> Constraint
  constraintTypes func constraint = go constraint
    where
      go (c1 :&: c2)             = go c1 :&: go c2
      go (i :<=: t)              = i :<=: func t
      go (Share     t)           = Share     (func t)
      go (Drop      t)           = Drop      (func t)
      go (Escape    t)           = Escape    (func t)
      go (Exhausted t)           = Exhausted (func t)
      go (t1  :<  t2 )           = func t1 :< func t2
      go (t1  :=: t2 )           = func t1 :=: func t2
      go (Solved t)              = Solved $ func t
      go (UnboxedNoRecurse rp s) = UnboxedNoRecurse rp s
      go Sat                     = Sat
      go Unsat                   = Unsat
  
  
  -- | Given a function that acts on 'Type' values, produce a function
  --   that acts on the types inside an 'Expr'.
  exprTypes :: (Type -> Type) -> Expr -> Expr
  exprTypes func expr = go expr
   where
    go (TypeApp f  ts     ) = TypeApp f (map func ts)
    go (Sig     e  t      ) = Sig (go e) (func t)
    go (PrimOp  o  es     ) = PrimOp o (map go es)
    go (Con     n  e      ) = Con n (go e)
    go (Apply   e1 e2     ) = Apply (go e1) (go e2)
    go (Struct fs         ) = Struct (map (fmap go) fs)
    go (If  e e1 e2       ) = If (go e) (go e1) (go e2)
    go (Let v e1 e2       ) = Let v (go e1) (go e2)
    go (LetBang vs v e1 e2) = LetBang vs v (go e1) (go e2)
    go (Take r f v e1 e2  ) = Take r f v (go e1) (go e2)
    go (Put e1 f e2       ) = Put (go e1) f (go e2)
    go (Member e f        ) = Member (go e) f
    go (Case e k x e1 y e2) = Case (go e) k x (go e1) y (go e2)
    go (Esac e k x e1     ) = Esac (go e) k x (go e1)
    go e                    = e
  
  -- | Given a 'RW.Rewrite' on types, apply it over every subterm in a type, i.e. recursively applying
  --   the rewrite to every subterm.
  --
  --   If the rewrite succeeds on a subterm, the rewrite is not run again on the result. Therefore,
  --   this is guaranteed to terminate.
  --
  --   This could be accomplished with a datatype generics library but is overkill in this case I
  --   think.
  traverseType :: (RW.Rewrite Type) -> Type -> Type
  traverseType func ty = case RW.run func ty of
    Just t' -> t'
    Nothing -> case ty of
      Record n es s ->
        Record n (Row.mapEntries (entryTypes (traverseType func)) es) s
      AbsType n s ts -> AbsType n s (map (traverseType func) ts)
      Variant es -> Variant (Row.mapEntries (entryTypes (traverseType func)) es)
      Function t1 t2 -> Function (traverseType func t1) (traverseType func t2)
      Bang t         -> Bang (traverseType func t)
      _              -> ty
  
  -- | Given a 'RW.Rewrite' on types, apply it over every subterm in a type, i.e. recursively applying
  --   the rewrite to every subterm.
  --
  --   If the rewrite succeeds on a subterm, the rewrite *is* run again on the result. Therefore,
  --   the rewrite must be a reduction or this function will not terminate.
  --
  --   If this function terminates, the result is guaranteed not to contain any further reducible
  --   subterm.
  normaliseType :: (RW.Rewrite Type) -> Type -> Type
  normaliseType func ty =
    let t' = fromMaybe ty (RW.run func ty)
    in
      case t' of
        Record n es s ->
          Record n (Row.mapEntries (entryTypes (normaliseType func)) es) s
        AbsType n s ts -> AbsType n s (map (normaliseType func) ts)
        Variant es ->
          Variant (Row.mapEntries (entryTypes (normaliseType func)) es)
        Function t1 t2 ->
          Function (normaliseType func t1) (normaliseType func t2)
        Bang t -> Bang (normaliseType func t)
        _      -> t'
  
  
  -- | Checks if two recursive parameter bindings are the same type
  sameRecursive :: RecPar -> RecPar -> Bool
  sameRecursive (Rec _) (Rec _) = True
  sameRecursive None    None    = True
  sameRecursive _       _       = False
  
  -- | Unrolls a recursive parameter to the record it references
  unroll :: RecParName -> RecContext -> Type
  unroll n (Just ctxt) = mapRecPars (Just ctxt)  (ctxt M.! n)
  unroll _ _ = error "Impossible: cannot unroll a recursive parameter with an empty context"
  
  -- | Given a context, changes all recursive parameter references from TypeVar to RecPar according to the context
  mapRecPars :: RecContext -> Type -> Type
  mapRecPars ctxt (AbsType n s ts)   = AbsType n s $ map (mapRecPars ctxt) ts
  mapRecPars ctxt (Variant row)      = Variant $ Row.mapEntries (\(Entry n t tk) -> Entry n (mapRecPars ctxt t) tk) row
  mapRecPars ctxt (Bang t)           = Bang $ mapRecPars ctxt t
  mapRecPars ctxt (RecPar n Nothing    ) = RecPar n ctxt
  mapRecPars ctxt (RecParBang n Nothing) = RecParBang n ctxt
  mapRecPars ctxt (RecPar n _    ) = error "Impossible: mapRecPars found a recursive parameter with a context inside a context"
  mapRecPars ctxt (RecParBang n _) = error "Impossible: mapRecPars found a recursive parameter (banged) with a context inside a context"
  mapRecPars (Just ctxt) r@(Record par row s) = 
    Record par (Row.mapEntries (\(Entry n t tk) -> Entry n (mapRecPars (addRecPar par) t) tk) row) s
    where addRecPar p = Just $ case p of
                          Rec v -> (M.insert v r ctxt)
                          _ -> ctxt
  mapRecPars ctxt (Function a b) = Function (mapRecPars ctxt a) (mapRecPars ctxt b)
  mapRecPars _ t = t
  
  
  
  -- | A rewrite that substitutes a given unification type variable for a type term in a type.
  substUV :: (VarName, Type) -> RW.Rewrite Type
  substUV (x, t) = RW.rewrite $
    \t' -> case t' of
      (UnifVar v) | x == v -> Just t
      _                    -> Nothing
  
  -- | A rewrite that substitutes a given unification row variable for a row in a type.
  substRowV :: (VarName, Row) -> RW.Rewrite Type
  substRowV (x, (Row m' q)) = RW.rewrite $
    \t' -> case t' of
      Variant (Row m (Just v)) | x == v -> Just (Variant (Row (M.union m m') q))
      Record n (Row m (Just v)) s | x == v ->
        Just (Record n (Row (M.union m m') q) s)
      _ -> Nothing
  
  -- | A rewrite that substitutes a given unification sigil variable for a sigil in a type.
  substSigilV :: (VarName, Sigil) -> RW.Rewrite Type
  substSigilV (x, s) = RW.rewrite $
    \t' -> case t' of
      Record n r (UnknownSigil v) | x == v -> Just (Record n r s)
      _ -> Nothing
  
  -- | A rewrite that substitutes a rigid type variable for a type term in a type.
  substTV :: (VarName, Type) -> RW.Rewrite Type
  substTV (x, t) = RW.rewrite $
    \t' -> case t' of
      (TypeVar v) | x == v     -> Just t
      (TypeVarBang v) | x == v -> Just (Bang t)
      _                        -> Nothing
  
  -- | A rewrite that substitutes the unkown recursive parameter on a boxed record for a parameter
  substRecPar :: (VarName, RecPar) -> RW.Rewrite Type
  substRecPar (v1, v2) = RW.rewrite $
    \t' -> case t' of
      Record (UnknownParameter n) r s | n == v1 -> 
        Just (Record v2 r s)
      _ -> Nothing
  
  -- | A convenience that allows multiple substitutions to type variables to be made simulatenously.
  substTVs :: [(VarName, Type)] -> RW.Rewrite Type
  substTVs = foldMap substTV
  
  -- | A convenience that allows multiple substitutions to unification type variables to be made
  --   simulatenously.
  substUVs :: [(VarName, Type)] -> RW.Rewrite Type
  substUVs = foldMap substUV
  
  -- | Just 'traverseType' composed with 'substTVs'
  traverseSubstTVs :: [(VarName, Type)] -> Type -> Type
  traverseSubstTVs = traverseType . substTVs
  
  
  -- | Returns @True@ iff the given integer fits within the given primitive type without overflow.
  fits :: Integer -> PrimType -> Bool
  fits i U8  = i >= 0 && i <= 255
  fits i U16 = i >= 0 && i <= 65535
  fits i U32 = i >= 0 && i <= 4294967295
  fits i U64 = i >= 0 && i <= 18446744073709551615
  fits _ _   = False
  
  -- | Returns @True@ if the two inputs are equal, or if either of them are an unknown sigil
  --   variable (morally, in this case the two inputs could be made equal by unification).
  sigilsCompatible :: Sigil -> Sigil -> Bool
  sigilsCompatible (UnknownSigil{}) y                = True
  sigilsCompatible x                (UnknownSigil{}) = True
  sigilsCompatible x                y                = x == y
  
  -- | Run a 'Fresh' computation with 'unifVars' as the source of fresh names.
  withUnifVars :: Monad m => FreshT VarName m a -> m a
  withUnifVars fr = evalFreshT unifVars fr
  
  -- | A stream of greek unification variable names.
  unifVars :: S.Stream VarName
  unifVars = S.fromList names
    where
      names = [ g:n | n <- nums, g <- "𝛂𝛃𝛄𝛅𝛆𝛇𝛈𝛉𝛊𝛋𝛍𝛎𝛏𝛑𝛖𝛗𝛘𝛙" ]
      nums = "":map show [1 :: Integer ..]