{-# OPTIONS_GHC -Wall #-}
{-# LANGUAGE OverloadedStrings #-}
module Nitpick.PatternMatches (patternMatches) where

{- The algorithm used here comes from "Warnings for Pattern Matching"
by Luc Maranget. Check it out for more information!

http://moscova.inria.fr/~maranget/papers/warn/warn.pdf

-}

import Control.Arrow ((***), second)
import qualified Data.Foldable as F
import qualified Data.List as List
import qualified Data.Map as Map
import qualified Data.Maybe as Maybe
import qualified Data.Set as Set
import qualified Data.Text as Text
import Data.Text (Text)

import qualified AST.Expression.Canonical as C
import qualified AST.Helpers as Help
import qualified AST.Literal as L
import qualified AST.Module as Module
import qualified AST.Module.Name as ModuleName
import qualified AST.Pattern as Pattern
import qualified AST.Variable as Var
import Elm.Utils ((|>))
import Nitpick.Pattern (Pattern(..), fromCanonicalPattern)
import qualified Optimize.DecisionTree as DT
import qualified Reporting.Annotation as A
import qualified Reporting.Error.Pattern as Error
import qualified Reporting.Region as Region
import qualified Reporting.Result as Result



-- NITPICK PATTERN MATCHES


type Result warning a =
  Result.Result () warning Error.Error a


patternMatches :: Module.Interfaces -> Module.Canonical -> Result wrn DT.VariantDict
patternMatches interfaces (Module.Module name info) =
  let
    arityDict =
      toArityDict interfaces name (Module.unions info)

    variantDict =
      Map.map (Map.map _size) arityDict
  in
    pure variantDict
      <* checkExpression arityDict (Module.program info)


checkExpression :: ArityDict -> C.Expr -> Result wrn ()
checkExpression arityDict (A.A region expression) =
  let
    go =
      checkExpression arityDict

    go2 a b =
      go a <* go b
  in
  case expression of
    C.Literal _ ->
        Result.ok ()

    C.Var _ ->
        Result.ok ()

    C.List listExprs ->
        F.traverse_ go listExprs

    C.Binop _ leftExpr rightExpr ->
        go2 leftExpr rightExpr

    C.Lambda pattern@(A.A patRegion _) body ->
        checkPatterns arityDict patRegion Error.Arg [pattern]
        <* go body

    C.App func arg ->
        go2 func arg

    C.If branches finally ->
        F.traverse_ (uncurry go2) branches
        <* go finally

    C.Let defs body ->
        go body
          <* F.traverse_ goDef defs
      where
        goDef (C.Def _ pattern@(A.A patRegion _) expr _) =
            checkPatterns arityDict patRegion Error.LetBound [pattern]
            <* go expr

    C.Case expr branches ->
        go expr
        <* checkPatterns arityDict region Error.Case (map fst branches)
        <* F.traverse_ (go . snd) branches

    C.Ctor _ctor exprs ->
        F.traverse_ go exprs

    C.Access record _field ->
        go record

    C.Update record fields ->
        go record
        <* F.traverse_ (go . snd) fields

    C.Record fields ->
        F.traverse_ (go . snd) fields

    C.Cmd _ _ ->
        Result.ok ()

    C.Sub _ _ ->
        Result.ok ()

    C.OutgoingPort _ _ ->
        Result.ok ()

    C.IncomingPort _ _ ->
        Result.ok ()

    C.Program _ _ ->
        error "DANGER - Program AST nodes should not be in Nitpick.PatternMatches."

    C.SaveEnv _ _ ->
        Result.ok ()

    C.GLShader _ _ _ ->
        Result.ok ()



-- CHECK PATTERNS


checkPatterns :: ArityDict -> Region.Region -> Error.Origin -> [Pattern.Canonical] -> Result wrn ()
checkPatterns arityDict region origin patterns =
  do  matrix <- checkRedundant arityDict region [] patterns
      case isExhaustive arityDict matrix 1 of
        [] ->
          Result.ok ()

        badPatterns ->
          Result.throw region (Error.Incomplete origin (map head badPatterns))



-- EXHAUSTIVE PATTERNS


isExhaustive :: ArityDict -> [[Pattern]] -> Int -> [[Pattern]]
isExhaustive arityDict matrix n =
  case (matrix, n) of
    ([], _) ->
      [replicate n Anything]

    (_, 0) ->
      []

    (_, _) ->
      let
        ctorSet =
          List.foldl' isCompleteHelp Set.empty matrix

        actual =
          Set.size ctorSet
      in
        if actual == 0 then
          (:) Anything
            <$> isExhaustive arityDict (Maybe.mapMaybe toDefault matrix) (n - 1)

        else
          let
            (ArityInfo expected info) =
              getArityInfo (Set.findMin ctorSet) arityDict
          in
            if actual < expected then
              (:)
                <$> Maybe.mapMaybe (isMissing ctorSet) info
                <*> isExhaustive arityDict (Maybe.mapMaybe toDefault matrix) (n - 1)

            else
              let
                isExhaustiveHelp (name, arity) =
                  recoverCtor name arity <$>
                  isExhaustive
                    arityDict
                    (Maybe.mapMaybe (specialize name arity) matrix)
                    (arity + n - 1)
              in
                concatMap isExhaustiveHelp info


isMissing :: Set.Set Var.Canonical -> (Var.Canonical, Int) -> Maybe Pattern
isMissing ctorSet (name, arity) =
  if Set.member name ctorSet then
    Nothing

  else
    Just (Ctor name (replicate arity Anything))


recoverCtor :: Var.Canonical -> Int -> [Pattern] -> [Pattern]
recoverCtor name arity patterns =
  let
    (args, rest) =
      splitAt arity patterns
  in
    Ctor name args : rest



-- REDUNDANT PATTERNS


checkRedundant
  :: ArityDict
  -> Region.Region
  -> [[Pattern]]
  -> [Pattern.Canonical]
  -> Result wrn [[Pattern]]
checkRedundant arityDict region checked unchecked =
  case unchecked of
    [] ->
      Result.ok checked

    nextPattern@(A.A patRegion _) : rest ->
      let
        next =
          [fromCanonicalPattern nextPattern]
      in
        if isUseful arityDict checked next then
          checkRedundant arityDict region (next : checked) rest
        else
          Result.throw region (Error.Redundant patRegion (length checked + 1))


isUseful :: ArityDict -> [[Pattern]] -> [Pattern] -> Bool
isUseful arityDict matrix vector =
  case (matrix, vector) of
    ([], _) ->
      True

    (_, []) ->
      False

    (_, firstPattern : patterns) ->
      case firstPattern of
        Ctor name args ->
          let
            newMatrix =
              Maybe.mapMaybe (specialize name (length args)) matrix
          in
            isUseful arityDict newMatrix (args ++ patterns)

        Anything ->
          case isComplete arityDict matrix of
            Nothing ->
              isUseful arityDict (Maybe.mapMaybe toDefault matrix) patterns

            Just arityInfo ->
              let
                isUsefulHelp (ctor, arity) =
                  isUseful
                    arityDict
                    (Maybe.mapMaybe (specialize ctor arity) matrix)
                    (replicate arity Anything ++ patterns)
              in
                any isUsefulHelp arityInfo

        Literal literal ->
          let
            newMatrix =
              Maybe.mapMaybe (specializeLiteral literal) matrix
          in
            isUseful arityDict newMatrix patterns



-- SPECIALIZE MATRICES


specialize :: Var.Canonical -> Int -> [Pattern] -> Maybe [Pattern]
specialize ctorName arity row =
  case row of
    [] ->
      error "Compiler error! Empty matrices should not get specialized."

    firstPattern : patterns ->
      case firstPattern of
        Ctor name args ->
          if name == ctorName then Just (args ++ patterns) else Nothing

        Anything ->
          Just (replicate arity Anything ++ patterns)

        Literal _ ->
          error $
            "Compiler bug! After type checking, constructors and literals\
            \ should never align in pattern match exhaustiveness checks."


specializeLiteral :: L.Literal -> [Pattern] -> Maybe [Pattern]
specializeLiteral literal row =
  case row of
    [] ->
      error "Compiler error! Empty matrices should not get specialized."

    firstPattern : patterns ->
      case firstPattern of
        Literal lit ->
          if lit == literal then Just patterns else Nothing

        Anything ->
          Just patterns

        Ctor _ _ ->
          error $
            "Compiler bug! After type checking, constructors and literals\
            \ should never align in pattern match exhaustiveness checks."


toDefault :: [Pattern] -> Maybe [Pattern]
toDefault row =
  case row of
    [] ->
      Nothing

    Ctor _ _ : _ ->
      Nothing

    Anything : patterns ->
      Just patterns

    Literal _ : _ ->
      Nothing



-- ALL CONSTRUCTORS ARE PRESENT?


isComplete :: ArityDict -> [[Pattern]] -> Maybe [(Var.Canonical, Int)]
isComplete arityDict matrix =
  let
    ctorSet =
      List.foldl' isCompleteHelp Set.empty matrix

    actual =
      Set.size ctorSet
  in
    if actual == 0 then
      Nothing

    else
      let
        (ArityInfo expected info) =
          getArityInfo (Set.findMin ctorSet) arityDict
      in
        if expected == actual then Just info else Nothing


isCompleteHelp :: Set.Set Var.Canonical -> [Pattern] -> Set.Set Var.Canonical
isCompleteHelp ctors row =
  case row of
    Ctor name _ : _ ->
      Set.insert name ctors

    _ ->
      ctors



-- ARITY DICT


type ArityDict =
  Map.Map Var.Home (Map.Map Text ArityInfo)


data ArityInfo =
  ArityInfo
    { _size :: !Int
    , _info :: [(Var.Canonical, Int)]
    }


toArityDict :: Module.Interfaces -> ModuleName.Canonical -> Module.Unions -> ArityDict
toArityDict interfaces localName localUnions =
  interfaces
    |> Map.mapWithKey (\name iface -> toArityEntry name (Module.iUnions iface))
    |> Map.mapKeysMonotonic Var.Module
    |> Map.insert (Var.Module localName) (toArityEntry localName localUnions)
    |> Map.union builtinDict


builtinDict :: ArityDict
builtinDict =
  let
    listInfo =
      ArityInfo 2 [ (Var.cons, 2), (Var.nil, 0) ]

    boolInfo =
      ArityInfo 2 [ (Var.true, 0), (Var.false, 0) ]
  in
    Map.singleton Var.BuiltIn $ Map.fromList $
      [ ("::", listInfo)
      , ("[]", listInfo)
      , ("True", boolInfo)
      , ("False", boolInfo)
      ]
      ++ map makeTupleInfo [0..8]


makeTupleInfo :: Int -> ( Text, ArityInfo )
makeTupleInfo n =
  ( Help.makeTuple n, ArityInfo 1 [ (Var.tuple n, n) ] )


toArityEntry :: ModuleName.Canonical -> Module.Unions -> Map.Map Text ArityInfo
toArityEntry name unions =
  Map.fromList (concatMap (toArityEntryHelp name) (Map.elems unions))


toArityEntryHelp :: ModuleName.Canonical -> Module.UnionInfo Text -> [(Text, ArityInfo)]
toArityEntryHelp name (_tvars, ctors) =
  let
    arityInfo =
      ArityInfo (length ctors) (map (Var.fromModule name *** length) ctors)
  in
    map (second (\_ -> arityInfo)) ctors


getArityInfo :: Var.Canonical -> ArityDict -> ArityInfo
getArityInfo var@(Var.Canonical home name) arityDict =
  case Map.lookup name =<< Map.lookup home arityDict of
    Just arityInfo ->
      arityInfo

    Nothing ->
      if Help.isTuple name then
        ArityInfo 1 [ (var, read (drop 6 (Text.unpack name))) ]
      else
        error
          "Since the Nitpick phase happens after canonicalization and type\
          \ inference, it is impossible that a pattern in a case cannot be\
          \ found."