Exhaustiveness.hs

{-# LANGUAGE PatternSynonyms #-}

-- |
-- Module      :  Disco.Exhaustiveness
-- Copyright   :  disco team and contributors
-- Maintainer  :  byorgey@gmail.com
--
-- SPDX-License-Identifier: BSD-3-Clause
--
-- Entry point into the exhaustiveness checker.
-- Converts information into a format the checker understands,
-- then pretty prints the results of running it.
module Disco.Exhaustiveness where

import Control.Monad (forM, unless, zipWithM)
import Control.Monad.Trans.Maybe (MaybeT, runMaybeT)
import Data.List (nub)
import Data.List.NonEmpty (NonEmpty)
import qualified Data.List.NonEmpty as NonEmpty
import Data.Maybe (catMaybes)
import Disco.AST.Generic (Side (..))
import Disco.AST.Surface (
  Pattern,
  prettyPatternP,
  pattern PBool,
  pattern PChar,
  pattern PInj,
  pattern PList,
  pattern PNat,
  pattern PNeg,
  pattern PString,
  pattern PTup,
  pattern PUnit,
  pattern PWild,
 )
import Disco.AST.Typed (
  APattern,
  ATerm,
  pattern APBool,
  pattern APChar,
  pattern APCons,
  pattern APInj,
  pattern APList,
  pattern APNat,
  pattern APNeg,
  pattern APString,
  pattern APTup,
  pattern APUnit,
  pattern APVar,
  pattern APWild,
 )
import Disco.Effects.Fresh (Fresh)
import Disco.Effects.LFresh (LFresh, runLFresh)
import qualified Disco.Exhaustiveness.Constraint as C
import qualified Disco.Exhaustiveness.Possibilities as Poss
import qualified Disco.Exhaustiveness.TypeInfo as TI
import Disco.Messages
import Disco.Pretty (Doc, initPA, withPA)
import qualified Disco.Pretty.DSL as DSL
import Disco.Pretty.Prec (PA, funPA)
import qualified Disco.Types as Ty
import Polysemy
import Polysemy.Output
import Polysemy.Reader
import Unbound.Generics.LocallyNameless (Name)

-- | This exhaustiveness checking algorithm is based on the paper
--   "Lower Your Guards: A Compositional Pattern-Match Coverage Checker":
--   https://www.microsoft.com/en-us/research/uploads/prod/2020/03/lyg.pdf
--
--   Some simplifications were made to adapt the algorithm to suit Disco.
--   The most notable change is that here we always generate (at most) 3
--   concrete examples of uncovered patterns, instead of finding the most
--   general complete description of every uncovered input.
checkClauses :: (Members '[Fresh, Reader Ty.TyDefCtx, Output (Message ann), Embed IO] r) => Name ATerm -> [Ty.Type] -> NonEmpty [APattern] -> Sem r ()
checkClauses name types pats = do
  args <- TI.newVars types
  cl <- zipWithM (desugarClause args) [1 ..] (NonEmpty.toList pats)
  let gdt = foldr1 Branch cl

  let argsNref = []
  (normalizedNrefs, _) <- ua [argsNref] gdt

  examples <- findPosExamples normalizedNrefs args

  unless (null examples) $ do
    let prettyExampleArgs exArgs =
          DSL.hcat $ map prettyPrintExample exArgs

    let prettyExampleLine prettyArgs =
          DSL.text (show name) DSL.<> prettyArgs DSL.<+> DSL.text "= ..."

    let prettyExamples =
          DSL.vcat $ map (prettyExampleLine . prettyExampleArgs) examples

    warn $
      DSL.text "Warning: the function"
        DSL.<+> DSL.text (show name)
        DSL.<+> DSL.text "is undefined for some inputs. For example:"
        DSL.$+$ prettyExamples

prettyPrintExample :: ExamplePat -> Sem r (Doc ann)
prettyPrintExample = runLFresh . runReader initPA . prettyPrintPattern . exampleToDiscoPattern

prettyPrintPattern :: (Members '[Reader PA, LFresh] r) => Pattern -> Sem r (Doc ann)
prettyPrintPattern = withPA funPA . prettyPatternP

exampleToDiscoPattern :: ExamplePat -> Pattern
exampleToDiscoPattern e@(ExamplePat TI.DataCon {TI.dcIdent = ident, TI.dcTypes = types} args) = case (ident, args) of
  (TI.KUnknown, _) -> PWild
  (TI.KUnit, _) -> PUnit
  (TI.KBool b, _) -> PBool b
  (TI.KNat n, _) -> PNat n
  (TI.KInt z, _) ->
    if z >= 0
      then PNat z
      else PNeg (PNat (abs z))
  (TI.KPair, _) -> PTup $ map exampleToDiscoPattern $ resugarPair e
  (TI.KCons, _) ->
    if take 1 types == [Ty.TyC]
      then PString $ resugarString e
      else PList $ map exampleToDiscoPattern $ resugarList e
  (TI.KNil, _) -> PList []
  (TI.KChar c, _) -> PChar c
  (TI.KLeft, [l]) -> PInj L $ exampleToDiscoPattern l
  (TI.KRight, [r]) -> PInj R $ exampleToDiscoPattern r
  (TI.KLeft, _) -> error "Found KLeft data constructor with 0 or multiple arguments"
  (TI.KRight, _) -> error "Found KRight data constructor with 0 or multiple arguments"

resugarPair :: ExamplePat -> [ExamplePat]
resugarPair e@(ExamplePat TI.DataCon {TI.dcIdent = ident} args) = case (ident, args) of
  (TI.KPair, [efst, esnd]) -> efst : resugarPair esnd
  (_, _) -> [e]

resugarList :: ExamplePat -> [ExamplePat]
resugarList (ExamplePat TI.DataCon {TI.dcIdent = ident} args) = case (ident, args) of
  (TI.KCons, [ehead, etail]) -> ehead : resugarList etail
  (_, _) -> []

resugarString :: ExamplePat -> String
resugarString (ExamplePat TI.DataCon {TI.dcIdent = ident} args) = case (ident, args) of
  (TI.KCons, [ehead, etail]) -> assumeExampleChar ehead : resugarString etail
  (_, _) -> []

assumeExampleChar :: ExamplePat -> Char
assumeExampleChar (ExamplePat TI.DataCon {TI.dcIdent = TI.KChar c} _) = c
assumeExampleChar _ = error "Wrongly assumed that an ExamplePat was a Char"

desugarClause :: (Members '[Fresh, Embed IO] r) => [TI.TypedVar] -> Int -> [APattern] -> Sem r Gdt
desugarClause args clauseIdx argsPats = do
  guards <- zipWithM desugarMatch args argsPats
  return $ foldr Guarded (Grhs clauseIdx) $ concat guards

-- To work with the LYG algorithm, we need to desugar n-tuples to nested pairs
-- Just having a Tuple type with a variable number of arguments breaks.
-- Imagine we have
-- foo (1,2,3) = ...
-- foo (1,(2,n)) = ...
-- if we keep things in our nice "sugared" form, the solver will get confused,
-- and not realize that the last element of the tuple is fully covered by n,
-- because there will be two notions of last element: the last in the triple and
-- the last in the nested pair
desugarTuplePats :: [APattern] -> APattern
desugarTuplePats [] = error "Found empty tuple, what happened?"
desugarTuplePats [p] = p
desugarTuplePats (pfst : rest) = APTup (Ty.getType pfst Ty.:*: Ty.getType psnd) [pfst, psnd]
 where
  psnd = desugarTuplePats rest

-- | Convert a Disco APattern into a list of Guards which cover that pattern
--
--   These patterns are currently not handled:
--     , APNeg     --still need to handle rational case
--     , APFrac    --required for rationals?
--     algebraic (probably will be eventually replaced anyway):
--     , APAdd
--     , APMul
--     , APSub
--   These (or some updated version of them) may be handled eventually,
--   once updated arithmetic patterns are merged.
--
--   We treat unhandled patterns as if they are exhaustively matched against
--   (aka, they are seen as wildcards by the checker).
--   This necessarily results in some false negatives, but no false positives.
desugarMatch :: (Members '[Fresh, Embed IO] r) => TI.TypedVar -> APattern -> Sem r [Guard]
desugarMatch var pat = case pat of
  (APTup (ta Ty.:*: tb) [pfst, psnd]) -> do
    varFst <- TI.newVar ta
    varSnd <- TI.newVar tb
    guardsFst <- desugarMatch varFst pfst
    guardsSnd <- desugarMatch varSnd psnd
    let guardPair = (var, GMatch (TI.pair ta tb) [varFst, varSnd])
    return $ [guardPair] ++ guardsFst ++ guardsSnd
  (APTup _ sugary) -> desugarMatch var (desugarTuplePats sugary)
  (APCons _ subHead subTail) -> do
    let typeHead = Ty.getType subHead
    let typeTail = Ty.getType subTail
    varHead <- TI.newVar typeHead
    varTail <- TI.newVar typeTail
    guardsHead <- desugarMatch varHead subHead
    guardsTail <- desugarMatch varTail subTail
    let guardCons = (var, GMatch (TI.cons typeHead typeTail) [varHead, varTail])
    return $ [guardCons] ++ guardsHead ++ guardsTail
  -- We have to desugar Lists into Cons and Nils
  (APList _ []) -> return [(var, GMatch TI.nil [])]
  (APList ty (phead : ptail)) -> do
    -- APCons have the type of the list they are part of
    desugarMatch var (APCons ty phead (APList ty ptail))
  -- we have to desugar to a list, becuse we can match strings with cons
  (APString str) -> do
    let strType = Ty.TyList Ty.TyC
    desugarMatch var (APList strType (map APChar str))
  -- A bit of strangeness is required here because of how patterns work
  (APNat Ty.TyN nat) -> return [(var, GMatch (TI.natural nat) [])]
  (APNat Ty.TyZ z) -> return [(var, GMatch (TI.integer z) [])]
  (APNeg Ty.TyZ (APNat Ty.TyN z)) -> return [(var, GMatch (TI.integer (-z)) [])]
  -- These are more straightforward:
  (APWild _) -> return []
  (APVar ty name) -> do
    let newAlias = TI.TypedVar (name, ty)
    return [(newAlias, GWasOriginally var)]
  APUnit -> return [(var, GMatch TI.unit [])]
  (APBool b) -> return [(var, GMatch (TI.bool b) [])]
  (APChar c) -> return [(var, GMatch (TI.char c) [])]
  (APInj (tl Ty.:+: tr) side subPat) -> do
    let dc = case side of
          L -> TI.left tl
          R -> TI.right tr
    newVar <- case side of
      L -> TI.newVar tl
      R -> TI.newVar tr
    guards <- desugarMatch newVar subPat
    return $ (var, GMatch dc [newVar]) : guards
  _ -> return []

data Gdt where
  Grhs :: Int -> Gdt
  Branch :: Gdt -> Gdt -> Gdt
  Guarded :: Guard -> Gdt -> Gdt
  deriving (Show, Eq)

type Guard = (TI.TypedVar, GuardConstraint)

data GuardConstraint where
  GMatch :: TI.DataCon -> [TI.TypedVar] -> GuardConstraint
  GWasOriginally :: TI.TypedVar -> GuardConstraint
  deriving (Show, Eq)

newtype Literal = Literal (TI.TypedVar, LitCond)

data LitCond where
  LitMatch :: TI.DataCon -> [TI.TypedVar] -> LitCond
  LitNot :: TI.DataCon -> LitCond
  LitWasOriginally :: TI.TypedVar -> LitCond
  deriving (Show, Eq, Ord)

data Ant where
  AGrhs :: [C.NormRefType] -> Int -> Ant
  ABranch :: Ant -> Ant -> Ant
  deriving (Show)

ua :: (Members '[Fresh, Reader Ty.TyDefCtx] r) => [C.NormRefType] -> Gdt -> Sem r ([C.NormRefType], Ant)
ua nrefs gdt = case gdt of
  Grhs k -> return ([], AGrhs nrefs k)
  Branch t1 t2 -> do
    (n1, u1) <- ua nrefs t1
    (n2, u2) <- ua n1 t2
    return (n2, ABranch u1 u2)
  Guarded (y, GWasOriginally x) t -> do
    n <- addLitMulti nrefs $ Literal (y, LitWasOriginally x)
    ua n t
  Guarded (x, GMatch dc args) t -> do
    n <- addLitMulti nrefs $ Literal (x, LitMatch dc args)
    (n', u) <- ua n t
    n'' <- addLitMulti nrefs $ Literal (x, LitNot dc)
    return (n'' ++ n', u)

addLitMulti :: (Members '[Fresh, Reader Ty.TyDefCtx] r) => [C.NormRefType] -> Literal -> Sem r [C.NormRefType]
addLitMulti [] _ = return []
addLitMulti (n : ns) lit = do
  r <- runMaybeT $ addLiteral n lit
  case r of
    Nothing -> addLitMulti ns lit
    Just cfs -> do
      ns' <- addLitMulti ns lit
      return $ cfs : ns'

addLiteral :: (Members '[Fresh, Reader Ty.TyDefCtx] r) => C.NormRefType -> Literal -> MaybeT (Sem r) C.NormRefType
addLiteral constraints (Literal (x, c)) = case c of
  LitWasOriginally z ->
    constraints `C.addConstraint` (x, C.CWasOriginally z)
  LitMatch dc args ->
    constraints `C.addConstraint` (x, C.CMatch dc args)
  LitNot dc ->
    constraints `C.addConstraint` (x, C.CNot dc)

data InhabPat where
  IPIs :: TI.DataCon -> [InhabPat] -> InhabPat
  IPNot :: [TI.DataCon] -> InhabPat
  deriving (Show, Eq, Ord)

-- Sanity check: are we giving the dataconstructor the
-- correct number of arguments?
mkIPMatch :: TI.DataCon -> [InhabPat] -> InhabPat
mkIPMatch k pats =
  if length (TI.dcTypes k) /= length pats
    then error $ "Wrong number of DataCon args" ++ show (k, pats)
    else IPIs k pats

findInhabitants :: (Members '[Fresh, Reader Ty.TyDefCtx] r) => [C.NormRefType] -> [TI.TypedVar] -> Sem r (Poss.Possibilities [InhabPat])
findInhabitants nrefs args = do
  a <- forM nrefs (`findAllForNref` args)
  return $ Poss.anyOf a

findAllForNref :: (Members '[Fresh, Reader Ty.TyDefCtx] r) => C.NormRefType -> [TI.TypedVar] -> Sem r (Poss.Possibilities [InhabPat])
findAllForNref nref args = do
  argPats <- forM args (`findVarInhabitants` nref)
  return $ Poss.allCombinations argPats

findVarInhabitants :: (Members '[Fresh, Reader Ty.TyDefCtx] r) => TI.TypedVar -> C.NormRefType -> Sem r (Poss.Possibilities InhabPat)
findVarInhabitants var cns =
  case posMatch of
    Just (k, args) -> do
      argPossibilities <- findAllForNref cns args
      return (mkIPMatch k <$> argPossibilities)
    Nothing -> case nub negMatches of
      [] -> Poss.retSingle $ IPNot []
      neg -> do
        tyCtx <- ask @Ty.TyDefCtx
        case TI.tyDataCons (TI.getType var) tyCtx of
          TI.Infinite _ -> Poss.retSingle $ IPNot neg
          TI.Finite dcs ->
            do
              let tryAddDc dc = do
                    vars <- TI.newVars (TI.dcTypes dc)
                    runMaybeT (cns `C.addConstraint` (var, C.CMatch dc vars))

              -- Try to add a positive constraint for each data constructor
              -- to the current nref
              -- If any of these additions succeed, save that nref,
              -- it now has positive information
              posNrefs <- catMaybes <$> forM dcs tryAddDc

              if null posNrefs
                then Poss.retSingle $ IPNot []
                else Poss.anyOf <$> forM posNrefs (findVarInhabitants var)
 where
  constraintsOnX = C.onVar var cns
  posMatch = C.posMatch constraintsOnX
  negMatches = C.negMatches constraintsOnX

findRedundant :: (Members '[Fresh, Reader Ty.TyDefCtx] r) => Ant -> [TI.TypedVar] -> Sem r [Int]
findRedundant ant args = case ant of
  AGrhs ref i -> do
    uninhabited <- Poss.none <$> findInhabitants ref args
    return [i | uninhabited]
  ABranch a1 a2 -> mappend <$> findRedundant a1 args <*> findRedundant a2 args

data ExamplePat where
  ExamplePat :: TI.DataCon -> [ExamplePat] -> ExamplePat
  deriving (Show)

-- | Less general version of the above inhabitant finding function
--   returns a maximum of 3 possible args lists that haven't been matched against,
--   as to not overwhelm new users of the language.
--   This is essentially a DFS, and it has a bad habit of
--   trying to build infinite lists whenever it can, so we give it a max depth of 32
--   If we reach 32 levels of nested dataconstructors in this language,
--   it is pretty safe to assume we were chasing after an infinite structure
findPosExamples :: (Members '[Fresh, Reader Ty.TyDefCtx] r) => [C.NormRefType] -> [TI.TypedVar] -> Sem r [[ExamplePat]]
findPosExamples nrefs args = do
  a <- forM nrefs (\nref -> findAllPosForNref 32 nref args)
  return $ take 3 $ Poss.getPossibilities $ Poss.anyOf a

findAllPosForNref :: (Members '[Fresh, Reader Ty.TyDefCtx] r) => Int -> C.NormRefType -> [TI.TypedVar] -> Sem r (Poss.Possibilities [ExamplePat])
findAllPosForNref fuel nref args = do
  argPats <- forM args (\arg -> findVarPosExamples (fuel - 1) arg nref)
  return $ Poss.allCombinations argPats

findVarPosExamples :: (Members '[Fresh, Reader Ty.TyDefCtx] r) => Int -> TI.TypedVar -> C.NormRefType -> Sem r (Poss.Possibilities ExamplePat)
findVarPosExamples fuel var cns =
  if fuel < 0
    then return mempty
    else case posMatch of
      Just (k, args) -> do
        argPossibilities <- findAllPosForNref (fuel - 1) cns args
        return (mkExampleMatch k <$> argPossibilities)
      Nothing -> do
        tyCtx <- ask @Ty.TyDefCtx
        let dcs = getPosFrom (TI.getType var) tyCtx negMatches
        let tryAddDc dc = do
              vars <- TI.newVars (TI.dcTypes dc)
              runMaybeT (cns `C.addConstraint` (var, C.CMatch dc vars))
        -- Try to add a positive constraint for each data constructor
        -- to the current nref
        -- If any of these additions succeed, save that nref,
        -- it now has positive information
        posNrefs <- catMaybes <$> forM dcs tryAddDc

        Poss.anyOf <$> forM posNrefs (findVarPosExamples (fuel - 1) var)
 where
  constraintsOnX = C.onVar var cns
  posMatch = C.posMatch constraintsOnX
  negMatches = C.negMatches constraintsOnX

getPosFrom :: Ty.Type -> Ty.TyDefCtx -> [TI.DataCon] -> [TI.DataCon]
getPosFrom ty ctx neg = take 3 $ filter (`notElem` neg) dcs
 where
  dcs = case TI.tyDataCons ty ctx of
    TI.Finite dcs' -> dcs'
    TI.Infinite dcs' -> dcs'

mkExampleMatch :: TI.DataCon -> [ExamplePat] -> ExamplePat
mkExampleMatch k pats =
  if length (TI.dcTypes k) /= length pats
    then error $ "Wrong number of DataCon args" ++ show (k, pats)
    else ExamplePat k pats