Constraint.hs

-- |
-- Module      :  Disco.Exhaustiveness.Constraint
-- Copyright   :  disco team and contributors
-- Maintainer  :  byorgey@gmail.com
--
-- SPDX-License-Identifier: BSD-3-Clause
--
-- The heart of the Lower Your Guards algorithm.
-- Functions for converting constraints detailing
-- pattern matches into a normalized description
-- of everything that is left uncovered.
module Disco.Exhaustiveness.Constraint where

import Control.Applicative (Alternative)
import Control.Monad (foldM, guard)
import Control.Monad.Trans (lift)
import Control.Monad.Trans.Maybe (MaybeT, runMaybeT)
import Data.Bifunctor (first)
import Data.List (partition)
import Data.Maybe (isJust, listToMaybe, mapMaybe)
import Disco.Effects.Fresh (Fresh)
import qualified Disco.Exhaustiveness.TypeInfo as TI
import qualified Disco.Types as Ty
import Polysemy
import Polysemy.Reader

data Constraint where
  CMatch :: TI.DataCon -> [TI.TypedVar] -> Constraint
  CNot :: TI.DataCon -> Constraint
  CWasOriginally :: TI.TypedVar -> Constraint
  deriving (Show, Eq, Ord)

posMatch :: [Constraint] -> Maybe (TI.DataCon, [TI.TypedVar])
posMatch constraints = listToMaybe $ mapMaybe (\case (CMatch k ys) -> Just (k, ys); _ -> Nothing) constraints

negMatches :: [Constraint] -> [TI.DataCon]
negMatches = mapMaybe (\case (CNot k) -> Just k; _ -> Nothing)

type ConstraintFor = (TI.TypedVar, Constraint)

-- Resolves term equalities, finding the leftmost id for a variable
-- I believe I3 of section 3.4 allows us to
-- do a linear scan from right to left
lookupVar :: TI.TypedVar -> [ConstraintFor] -> TI.TypedVar
lookupVar x = foldr getNextId x
 where
  getNextId (x', CWasOriginally y) | x' == x = const y
  getNextId _ = id

alistLookup :: (Eq a) => a -> [(a, b)] -> [b]
alistLookup a = map snd . filter ((== a) . fst)

onVar :: TI.TypedVar -> [ConstraintFor] -> [Constraint]
onVar x cs = alistLookup (lookupVar x cs) cs

type NormRefType = [ConstraintFor]

addConstraints :: Members '[Fresh, Reader Ty.TyDefCtx] r => NormRefType -> [ConstraintFor] -> MaybeT (Sem r) NormRefType
addConstraints = foldM addConstraint

addConstraint :: Members '[Fresh, Reader Ty.TyDefCtx] r => NormRefType -> ConstraintFor -> MaybeT (Sem r) NormRefType
addConstraint cns (x, c) = do
  breakIf $ any (conflictsWith c) (onVar x cns)
  addConstraintHelper cns (lookupVar x cns, c)

addConstraintHelper :: Members '[Fresh, Reader Ty.TyDefCtx] r => NormRefType -> ConstraintFor -> MaybeT (Sem r) NormRefType
addConstraintHelper cns cf@(origX, c) = case c of
  --- Equation (10)
  CMatch k args -> do
    case getConstructorArgs k (onVar origX cns) of
      -- 10c
      Just args' ->
        addConstraints
          cns
          (zipWith (\a b -> (a, CWasOriginally b)) args args')
      Nothing -> return added
  --- Equation (11)
  CNot _ -> do
    inh <- lift (inhabited added origX)
    guard inh -- ensure that origX is still inhabited, as per I2
    return added
  -- Equation (14)
  CWasOriginally y -> do
    let origY = lookupVar y cns
    if origX == origY
      then return cns
      else do
        let (noX', withX') = partition ((/= origX) . fst) cns
        addConstraints (noX' ++ [cf]) (substituteVarIDs origY origX withX')
 where
  added = cns ++ [cf]

-----
----- Helper functions for adding constraints:
-----

breakIf :: (Alternative f) => Bool -> f ()
breakIf = guard . not

-- | Returns a predicate that returns true if another
--   constraint conflicts with the one given.
--   This alone is not sufficient to test
--   if a constraint can be added, but it
--   filters out the easy negatives early on
conflictsWith :: Constraint -> (Constraint -> Bool)
conflictsWith c = case c of
  CMatch k _ -> \case
    (CMatch k' _) | k /= k' -> True -- 10a
    (CNot k') | k == k' -> True -- 10b
    _ -> False
  CNot k -> \case
    (CMatch k' _) | k == k' -> True -- 11a
    _ -> False
  CWasOriginally _ -> const False

-- | Search for a MatchDataCon that is matching on k specifically
--   (there should be at most one, see I4 in section 3.4)
--   and if it exists, return the variable ids of its arguments
getConstructorArgs :: TI.DataCon -> [Constraint] -> Maybe [TI.TypedVar]
getConstructorArgs k cfs =
  listToMaybe $
    mapMaybe (\case (CMatch k' vs) | k' == k -> Just vs; _ -> Nothing) cfs

-- | substituting y *for* x
--   ie replace the second with the first, replace x with y
substituteVarIDs :: TI.TypedVar -> TI.TypedVar -> [ConstraintFor] -> [ConstraintFor]
substituteVarIDs y x = map (first subst)
 where
  subst var = if var == x then y else x

-- | Deals with I2 from section 3.4
--   if a variable in the context has a resolvable type, there must be at least one constructor
--   which can be instantiated without contradiction of the refinement type
--   This function tests if this is true
inhabited :: Members '[Fresh, Reader Ty.TyDefCtx] r => NormRefType -> TI.TypedVar -> Sem r Bool
inhabited n var = do
  tyCtx <- ask @Ty.TyDefCtx
  case TI.tyDataCons (TI.getType var) tyCtx of
    TI.Infinite _ -> return True -- assume opaque types are inhabited
    TI.Finite constructors -> do
      or <$> mapM (instantiate n var) constructors

-- | Attempts to "instantiate" a match of the dataconstructor k on x
--   If we can add the MatchDataCon constraint to the normalized refinement
--   type without contradiction (a Nothing value),
--   then x is inhabited by k and we return true
instantiate :: Members '[Fresh, Reader Ty.TyDefCtx] r => NormRefType -> TI.TypedVar -> TI.DataCon -> Sem r Bool
instantiate cns var k = do
  args <- TI.newVars $ TI.dcTypes k
  let attempt = cns `addConstraint` (var, CMatch k args)
  isJust <$> runMaybeT attempt