Rules.hs

-- |
-- 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