content: -- |
  -- Module      :  Disco.Types.Rules
  -- Copyright   :  disco team and contributors
  -- Maintainer  :  byorgey@gmail.com
  --
  -- SPDX-License-Identifier: BSD-3-Clause
  --
  -- "Disco.Types.Rules" defines some generic rules about arity,
  -- subtyping, and sorts for disco base types.
  module Disco.Types.Rules (
    -- * Arity
    Variance (..),
    arity,
  
    -- * Qualifiers
    Qualifier (..),
    bopQual,
  
    -- * Sorts
    Sort,
    topSort,
  
    -- * Subtyping rules
    Dir (..),
    other,
    isSubA,
    isSubB,
    isDirB,
    supertypes,
    subtypes,
    dirtypes,
  
    -- * Qualifier and sort rules
    hasQual,
    hasSort,
    qualRules,
    sortRules,
    pickSortBaseTy,
  )
  where
  
  import Control.Monad ((>=>))
  import Data.Foldable (Foldable (..))
  import Data.Map (Map)
  import qualified Data.Map as M
  import qualified Data.Set as S
  import Disco.Types
  import Disco.Types.Qualifiers
  import Prelude hiding (Foldable (..))
  
  ------------------------------------------------------------
  -- Arity
  ------------------------------------------------------------
  
  -- | A particular type argument can be either co- or contravariant
  --   with respect to subtyping.
  data Variance = Co | Contra
    deriving (Show, Read, Eq, Ord)
  
  -- | The arity of a type constructor is a list of variances,
  --   expressing both how many type arguments the constructor takes,
  --   and the variance of each argument.  This is used to decompose
  --   subtyping constraints.
  --
  --   For example, @arity CArr = [Contra, Co]@ since function arrow is
  --   contravariant in its first argument and covariant in its second.
  --   That is, @S1 -> T1 <: S2 -> T2@ (@<:@ means "is a subtype of") if
  --   and only if @S2 <: S1@ and @T1 <: T2@.
  arity :: Con -> [Variance]
  arity CArr = [Contra, Co]
  arity CProd = [Co, Co]
  arity CSum = [Co, Co]
  arity (CContainer _) = [Co]
  arity CMap = [Contra, Co]
  arity CGraph = [Co]
  arity (CUser _) = error "Impossible! arity CUser"
  
  -- CUsers should always be replaced by their definitions before arity
  -- is called.
  
  ------------------------------------------------------------
  -- Subtyping rules
  ------------------------------------------------------------
  
  -- | A "direction" for the subtyping relation (either subtype or
  --   supertype).
  data Dir = SubTy | SuperTy
    deriving (Eq, Ord, Read, Show)
  
  -- | Swap directions.
  other :: Dir -> Dir
  other SubTy = SuperTy
  other SuperTy = SubTy
  
  --------------------------------------------------
  -- Subtype checks
  
  -- | Check whether one atomic type is a subtype of the other. Returns
  --   @True@ if either they are equal, or if they are base types and
  --   'isSubB' returns true.
  isSubA :: Atom -> Atom -> Bool
  isSubA a1 a2 | a1 == a2 = True
  isSubA (ABase t1) (ABase t2) = isSubB t1 t2
  isSubA _ _ = False
  
  -- | Check whether one base type is a subtype of another.
  isSubB :: BaseTy -> BaseTy -> Bool
  isSubB b1 b2 | b1 == b2 = True
  isSubB N Z = True
  isSubB N F = True
  isSubB N Q = True
  isSubB Z Q = True
  isSubB F Q = True
  isSubB B P = True
  isSubB _ _ = False
  
  -- | Check whether one base type is a sub- or supertype of another.
  isDirB :: Dir -> BaseTy -> BaseTy -> Bool
  isDirB SubTy b1 b2 = isSubB b1 b2
  isDirB SuperTy b1 b2 = isSubB b2 b1
  
  -- | List all the supertypes of a given base type.
  supertypes :: BaseTy -> [BaseTy]
  supertypes N = [N, Z, F, Q]
  supertypes Z = [Z, Q]
  supertypes F = [F, Q]
  supertypes B = [B, P]
  supertypes ty = [ty]
  
  -- | List all the subtypes of a given base type.
  subtypes :: BaseTy -> [BaseTy]
  subtypes Q = [Q, F, Z, N]
  subtypes F = [F, N]
  subtypes Z = [Z, N]
  subtypes P = [P, B]
  subtypes ty = [ty]
  
  -- | List all the sub- or supertypes of a given base type.
  dirtypes :: Dir -> BaseTy -> [BaseTy]
  dirtypes SubTy = subtypes
  dirtypes SuperTy = supertypes
  
  ------------------------------------------------------------
  -- Qualifier and sort rules
  ------------------------------------------------------------
  
  -- | Check whether a given base type satisfies a qualifier.
  hasQual :: BaseTy -> Qualifier -> Bool
  hasQual P QCmp = False -- can't compare Props
  hasQual _ QCmp = True
  hasQual P QBasic = False
  hasQual _ QBasic = True
  hasQual P QSimple = False
  hasQual _ QSimple = True
  -- hasQual (Fin _) q     | q `elem` [QNum, QSub, QEnum] = True
  -- hasQual (Fin n) QDiv  = isPrime n
  hasQual b QNum = b `elem` [N, Z, F, Q]
  hasQual b QSub = b `elem` [Z, Q]
  hasQual b QDiv = b `elem` [F, Q]
  hasQual b QEnum = b `elem` [N, Z, F, Q, C]
  hasQual b QBool = b `elem` [B, P]
  
  -- | Check whether a base type has a certain sort, which simply
  --   amounts to whether it satisfies every qualifier in the sort.
  hasSort :: BaseTy -> Sort -> Bool
  hasSort = all . hasQual
  
  -- | 'qualRulesMap' encodes some of the rules by which applications of
  --   type constructors can satisfy various qualifiers.
  --
  --   Each constructor maps to a set of rules.  Each rule is a mapping
  --   from a qualifier to the list of qualifiers needed on the type
  --   constructor's arguments for the bigger type to satisfy the
  --   qualifier.
  --
  --   Note in Disco we can get away with any given qualifier requiring
  --   /at most one/ qualifier on each type argument.  Then we can
  --   derive the 'sortRules' by combining 'qualRules'.  In general,
  --   however, you could imagine some particular qualifier requiring a
  --   set of qualifiers (i.e. a general sort) on a type argument.  In
  --   that case one would just have to encode 'sortRules' directly.
  qualRulesMap :: Map Con (Map Qualifier [Maybe Qualifier])
  qualRulesMap =
    M.fromList
      [ CProd
          ==> M.fromList
            [ QCmp ==> [Just QCmp, Just QCmp]
            , QSimple ==> [Just QSimple, Just QSimple]
            ]
      , CSum
          ==> M.fromList
            [ QCmp ==> [Just QCmp, Just QCmp]
            , QSimple ==> [Just QSimple, Just QSimple]
            ]
      , CList
          ==> M.fromList
            [ QCmp ==> [Just QCmp]
            , QSimple ==> [Just QSimple]
            ]
      , CBag
          ==> M.fromList
            [ QCmp ==> [Just QCmp]
            , QSimple ==> [Just QSimple]
            ]
      , CSet
          ==> M.fromList
            [ QCmp ==> [Just QCmp]
            , QSimple ==> [Just QSimple]
            ]
      , CGraph
          ==> M.fromList
            [ QCmp ==> [Just QCmp]
            , QNum ==> [Nothing]
            ]
      , CMap
          ==> M.fromList
            [ QCmp ==> [Just QCmp, Just QCmp]
            ]
      ]
   where
    (==>) :: a -> b -> (a, b)
    (==>) = (,)
  
  -- We could (theoretically) make graphs and maps also be simple values if we require the map's values are also simple.
  
  -- Eventually we can easily imagine adding an opt-in mode where
  -- numeric operations can be used on pairs and functions, then the
  -- qualRules would become dependent on what language extension/mode
  -- was chosen.  For example we could have rules like
  --
  -- [ CArr ==> M.fromList
  --   [ QNum ==> [Nothing, Just QNum]  -- (a -> b) can be +, * iff b can
  --   , QSub ==> [Nothing, Just QSub]  -- ditto for subtraction
  --   , QDiv ==> [Nothing, Just QDiv]  -- and division
  --   ]
  -- , CProd ==> M.fromList
  --   [ QNum ==> [Just QNum, Just QNum] -- (a,b) can be +, * iff a and b can
  --   , QSub ==> [Just QSub, Just QSub] -- etc.
  --   , QDiv ==> [Just QDiv, Just QDiv]
  --   ]
  -- ]
  
  -- | Given a constructor T and a qualifier we want to hold of a type T
  --   t1 t2 ..., return a list of qualifiers that need to hold of t1,
  --   t2, ...
  qualRules :: Con -> Qualifier -> Maybe [Maybe Qualifier]
  -- T t1 t2 ... is basic (contains no Prop) iff t1, t2 ... all are.
  qualRules c QBasic = Just (map (const (Just QBasic)) (arity c))
  -- Otherwise, just look up in the qualRulesMap.
  qualRules c q = (M.lookup c >=> M.lookup q) qualRulesMap
  
  -- | @sortRules T s = [s1, ..., sn]@ means that sort @s@ holds of
  --   type @(T t1 ... tn)@ if and only if  @s1 t1 /\ ... /\ sn tn@.
  --   For now this is just derived directly from 'qualRules'.
  --
  --   This is the @arity@ function described in section 4.1 of Traytel et
  --   al.
  sortRules :: Con -> Sort -> Maybe [Sort]
  sortRules c s = do
    -- If any of the quals q in sort s are not in the map corresponding
    -- to tycon c, there's no way to make c an instance of q, so fail
    -- (the mapM will succeed only if all lookups succeed)
    needQuals <- mapM (qualRules c) (S.toList s)
  
    -- Otherwise we are left with a list (corresponding to all the quals
    -- in sort s) of lists (each one corresponds to the type args of c).
    -- We zip them together to produce a list of sorts.
    return $ foldl' (zipWith (\srt -> maybe srt (`S.insert` srt))) (repeat topSort) needQuals
  
  -- | Pick a base type (generally the "simplest") that satisfies a given sort.
  pickSortBaseTy :: Sort -> BaseTy
  pickSortBaseTy s
    | QDiv `S.member` s && QSub `S.member` s = Q
    | QDiv `S.member` s = F
    | QSub `S.member` s = Z
    | QNum `S.member` s = N
    | QCmp `S.member` s = N
    | QEnum `S.member` s = N
    | QBool `S.member` s = B
    | QSimple `S.member` s = N
    | otherwise = Unit
  