{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE BangPatterns #-}

-- | 'RoughMap' is an approximate finite map data structure keyed on
-- @['RoughMatchTc']@. This is useful when keying maps on lists of 'Type's
-- (e.g. an instance head).
module GHC.Core.RoughMap
  ( -- * RoughMatchTc
    RoughMatchTc(..)
  , isRoughWildcard
  , typeToRoughMatchTc
  , RoughMatchLookupTc(..)
  , typeToRoughMatchLookupTc
  , roughMatchTcToLookup
  , roughMatchTcs
  , roughMatchTcsLookup
  , instanceCantMatch

    -- * RoughMap
  , RoughMap
  , emptyRM
  , lookupRM
  , lookupRM'
  , insertRM
  , filterRM
  , filterMatchingRM
  , elemsRM
  , sizeRM
  , foldRM
  , unionRM
  ) where

import GHC.Prelude

import GHC.Data.Bag
import GHC.Core.TyCon
import GHC.Core.TyCo.Rep
import GHC.Core.Type
import GHC.Utils.Outputable
import GHC.Types.Name
import GHC.Types.Name.Env
import GHC.Builtin.Types.Prim( cONSTRAINTTyConName, tYPETyConName )

import Control.Monad (join)
import Data.Data (Data)
import GHC.Utils.Panic

{-
Note [RoughMap]
~~~~~~~~~~~~~~~
We often want to compute whether one type matches another. That is, given
`ty1` and `ty2`, we want to know whether `ty1` is a substitution instance of `ty2`.

We can bail out early by taking advantage of the following observation:

  If `ty2` is headed by a generative type constructor, say `tc`,
  but `ty1` is not headed by that same type constructor,
  then `ty1` does not match `ty2`.

The idea is that we can use a `RoughMap` as a pre-filter, to produce a
short-list of candidates to examine more closely.

This means we can avoid computing a full substitution if we represent types
as applications of known generative type constructors. So, after type synonym
expansion, we classify application heads into two categories ('RoughMatchTc')

  - `RM_KnownTc tc`: the head is the generative type constructor `tc`,
  - `RM_Wildcard`: anything else.

A (RoughMap val) is semantically a list of (key,[val]) pairs, where
   key :: [RoughMatchTc]
So, writing # for `OtherTc`, and Int for `KnownTc "Int"`, we might have
    [ ([#, Int, Maybe, #, Int], v1)
    , ([Int, #, List], v2 ]

This map is stored as a trie, so looking up a key is very fast.
See Note [Matching a RoughMap] and Note [Simple Matching Semantics] for details on
lookup.

We lookup a key of type [RoughMatchLookupTc], and return the list of all values whose
keys "match":

Given the above map, here are the results of some lookups:
   Lookup key       Result
   -------------------------
   [Int, Int]       [v1,v2] -- Matches because the prefix of both entries matches
   [Int,Int,List]   [v2]
   [Bool]           []

Notice that a single key can map to /multiple/ values.  E.g. if we started
with (Maybe Int, val1) and (Maybe Bool, val2), we'd generate a RoughMap
that is semantically the list   [( Maybe, [val1,val2] )]

Note [RoughMap and beta reduction]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
There is one tricky case we have to account for when matching a rough map due
to Note [Eta reduction for data families] in `GHC.Core.Coercion.Axiom`:
Consider that the user has written a program containing a data family:

> data family Fam a b
> data instance Fam Int a = SomeType  -- known henceforth as FamIntInst

The LHS of this instance will be eta reduced, as described in Note [Eta
reduction for data families]. Consequently, we will end up with a `FamInst`
with `fi_tcs = [KnownTc Int]`. Naturally, we need RoughMap to return this
instance when queried for an instance with template, e.g., `[KnownTc Fam,
KnownTc Int, KnownTc Char]`.

This explains the third clause of the mightMatch specification in Note [Simple Matching Semantics].
As soon as the lookup key runs out, the remaining instances might match.

This only matters for the data-family case of a FamInstEnv (see Note [Over-saturated matches]
in GHC.Core.FamInstEnv; it's irrelevantfor ClsInstEnv and for type-family instances.
But we use RoughMaps for all cases, so we are conservative.

Note [Matching a RoughMap]
~~~~~~~~~~~~~~~~~~~~~~~~~~
The /lookup key/ into a rough map (RoughMatchLookupTc) is slightly
different to the /insertion key/ (RoughMatchTc).  Like the insertion
key each lookup argument is classified to a simpler key which
describes what could match that position. There are three
possibilities:

* RML_KnownTc Name: The argument is headed by a known type
  constructor.  Example: 'Bool' is classified as 'RML_KnownTc Bool'
  and '[Int]' is classified as `RML_KnownTc []`

* RML_NoKnownTc: The argument is definitely not headed by any known
  type constructor.  Example: For instance matching 'a[sk], a[tau]' and 'F a[sk], F a[tau]'
  are classified as 'RML_NoKnownTc', for family instance matching no examples.

* RML_WildCard: The argument could match anything, we don't know
  enough about it. For instance matching no examples, for type family matching,
  things to do with variables.

The interesting case for instance matching is the second case, because it does not appear in
an insertion key. The second case arises in two situations:

1. The head of the application is a type variable. The type variable definitely
   doesn't match with any of the KnownTC instances so we can discard them all. For example:
    Show a[sk] or Show (a[sk] b[sk]). One place constraints like this arise is when
    typechecking derived instances.

2. The head of the application is a known type family.
   For example: F a[sk]. The application of F is stuck, and because
   F is a type family it won't match any KnownTC instance so it's safe to discard
   all these instances.

Of course, these two cases can still match instances of the form `forall a . Show a =>`,
and those instances are retained as they are classified as RM_WildCard instances.

Note [Matches vs Unifiers]
~~~~~~~~~~~~~~~~~~~~~~~~~~
The lookupRM' function returns a pair of potential /matches/ and potential /unifiers/.
The potential matches is likely to be much smaller than the bag of potential unifiers due
to the reasoning about rigid type variables described in Note [Matching a RoughMap].
On the other hand, the instances captured by the RML_NoKnownTC case can still potentially unify
with any instance (depending on the substitution of said rigid variable) so they can't be discounted
from the list of potential unifiers. This is achieved by the RML_NoKnownTC case continuing
the lookup for unifiers by replacing RML_NoKnownTC with RML_LookupOtherTC.

This distinction between matches and unifiers is also important for type families.
During normal type family lookup, we care about matches and when checking for consistency
we care about the unifiers. This is evident in the code as `lookup_fam_inst_env` is
parameterised over a lookup function which either performs matching checking or unification
checking.

In addition to this, we only care whether there are zero or non-zero potential
unifiers, even if we have many candidates, the search can stop before consulting
each candidate. We only need the full list of unifiers when displaying error messages.
Therefore the list is computed lazily so much work can be avoided constructing the
list in the first place.

Note [Simple Matching Semantics]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Suppose `rm` is a RoughMap representing a set of (key,vals) pairs,
  where key::[RoughMapTc] and val::a.
Suppose I look up a key lk :: [RoughMapLookupTc] in `rm`
Then I get back (matches, unifiers) where
   matches  = [ vals | (key,vals) <- rm, key `mightMatch` lk ]
   unifiers = [ vals | (key,vals) <- rm, key `mightUnify` lk ]

Where mightMatch is defined like this:

  mightMatch :: [RoughMapTc] -> [RoughMapLookupTc] -> Bool
  mightMatch []  []    = True   -- A perfectly sized match might match
  mightMatch key []    = True   -- A shorter lookup key matches everything
  mightMatch []  (_:_) = True   -- If the lookup key is longer, then still might match
                                -- Note [RoughMap and beta reduction]
  mightMatch (k:ks) (lk:lks) =
    = case (k,lk) of
         -- Standard case, matching on a specific known TyCon.
         (RM_KnownTc n1, RML_KnownTc n2) -> n1==n2 && mightMatch ks lks
         -- For example, if the key for 'Show Bool' is [RM_KnownTc Show, RM_KnownTc Bool]
         ---and we match against (Show a[sk]) [RM_KnownTc Show, RML_NoKnownTc]
         -- then Show Bool can never match Show a[sk] so return False.
         (RM_KnownTc _, RML_NoKnownTc)   -> False
         -- Wildcard cases don't inform us anything about the match.
         (RM_WildCard, _ )    -> mightMatch ks lks
         (_, RML_WildCard)    -> mightMatch ks lks

  -- Might unify is very similar to mightMatch apart from RML_NoKnownTc may
  -- unify with any instance.
  mightUnify :: [RoughMapTc] -> [RoughMapLookupTc] -> Bool
  mightUnify []  []    = True   -- A perfectly sized match might unify
  mightUnify key []    = True   -- A shorter lookup key matches everything
  mightUnify []  (_:_) = True
  mightUnify (k:ks) (lk:lks) =
    = case (k,lk) of
         (RM_KnownTc n1, RML_KnownTc n2) -> n1==n2 && mightUnify ks lks
         (RM_KnownTc _, RML_NoKnownTc)   -> mightUnify (k:ks) (RML_WildCard:lks)
         (RM_WildCard, _ )    -> mightUnify ks lks
         (_, RML_WildCard)    -> mightUnify ks lks


The guarantee that RoughMap provides is that

if
   insert_ty `tcMatchTy` lookup_ty
then definitely
   typeToRoughMatchTc insert_ty `mightMatch` typeToRoughMatchLookupTc lookup_ty
but not vice versa

this statement encodes the intuition that the RoughMap is used as a quick pre-filter
to remove instances from the matching pool. The contrapositive states that if the
RoughMap reports that the instance doesn't match then `tcMatchTy` will report that the
types don't match as well.

-}

{- *********************************************************************
*                                                                      *
                Rough matching
*                                                                      *
********************************************************************* -}

{- Note [Rough matching in class and family instances]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Consider
  instance C (Maybe [Tree a]) Bool
and suppose we are looking up
     C Bool Bool

We can very quickly rule the instance out, because the first
argument is headed by Maybe, whereas in the constraint we are looking
up has first argument headed by Bool.  These "headed by" TyCons are
called the "rough match TyCons" of the constraint or instance.
They are used for a quick filter, to check when an instance cannot
possibly match.

The main motivation is to avoid sucking in whole instance
declarations that are utterly useless.  See GHC.Core.InstEnv
Note [ClsInst laziness and the rough-match fields].

INVARIANT: a rough-match TyCons `tc` is always a real, generative tycon,
like Maybe or Either, including a newtype or a data family, both of
which are generative. It replies True to `isGenerativeTyCon tc Nominal`.

But it is never
    - A type synonym
      E.g. Int and (S Bool) might match
           if (S Bool) is a synonym for Int

    - A type family (#19336)
      E.g.   (Just a) and (F a) might match if (F a) reduces to (Just a)
             albeit perhaps only after 'a' is instantiated.
-}


-- Key for insertion into a RoughMap
data RoughMatchTc
  = RM_KnownTc Name  -- INVARIANT: Name refers to a TyCon tc that responds
                     -- true to `isGenerativeTyCon tc Nominal`. See
                     -- Note [Rough matching in class and family instances]
  | RM_WildCard      -- e.g. type variable at the head
  deriving( Data )

-- Key for lookup into a RoughMap
-- See Note [Matching a RoughMap]
data RoughMatchLookupTc
  = RML_KnownTc Name -- ^ The position only matches the specified KnownTc
  | RML_NoKnownTc    -- ^ The position definitely doesn't match any KnownTc
  | RML_WildCard     -- ^ The position can match anything
  deriving ( Data )

instance Outputable RoughMatchLookupTc where
    ppr (RML_KnownTc nm) = text "RML_KnownTc" <+> ppr nm
    ppr RML_NoKnownTc = text "RML_NoKnownTC"
    ppr RML_WildCard = text "_"

instance Outputable RoughMatchTc where
    ppr (RM_KnownTc nm) = text "KnownTc" <+> ppr nm
    ppr RM_WildCard = text "OtherTc"

instanceCantMatch :: [RoughMatchTc] -> [RoughMatchTc] -> Bool
-- (instanceCantMatch tcs1 tcs2) returns True if tcs1 cannot
-- possibly be instantiated to actual, nor vice versa;
-- False is non-committal
instanceCantMatch (mt : ts) (ma : as) = itemCantMatch mt ma || instanceCantMatch ts as
instanceCantMatch _         _         =  False  -- Safe

itemCantMatch :: RoughMatchTc -> RoughMatchTc -> Bool
itemCantMatch (RM_KnownTc t) (RM_KnownTc a) = t /= a
itemCantMatch _              _              = False

roughMatchTcToLookup :: RoughMatchTc -> RoughMatchLookupTc
roughMatchTcToLookup (RM_KnownTc n) = RML_KnownTc n
roughMatchTcToLookup RM_WildCard = RML_WildCard

isRoughWildcard :: RoughMatchTc -> Bool
isRoughWildcard RM_WildCard  = True
isRoughWildcard (RM_KnownTc {}) = False

roughMatchTcs :: [Type] -> [RoughMatchTc]
roughMatchTcs tys = map typeToRoughMatchTc tys

roughMatchTcsLookup :: [Type] -> [RoughMatchLookupTc]
roughMatchTcsLookup tys = map typeToRoughMatchLookupTc tys

typeToRoughMatchLookupTc :: Type -> RoughMatchLookupTc
typeToRoughMatchLookupTc ty
  -- Expand synonyms first, as explained in Note [Rough matching in class and family instances].
  -- Failing to do so led to #22985.
  | Just ty' <- coreView ty
  = typeToRoughMatchLookupTc ty'
  | CastTy ty' _ <- ty
  = typeToRoughMatchLookupTc ty'
  | otherwise
  = case splitAppTys ty of
        -- Case 1: Head of application is a type variable, does not match any KnownTc.
        (TyVarTy {}, _) -> RML_NoKnownTc

        (TyConApp tc _, _)
          -- Case 2: Head of application is a known type constructor, hence KnownTc.
          | not (isTypeFamilyTyCon tc) -> RML_KnownTc $! roughMatchTyConName tc
          -- Case 3: Head is a type family so it's stuck and therefore doesn't match
          -- any KnownTc
          | isTypeFamilyTyCon tc -> RML_NoKnownTc

        -- Fallthrough: Otherwise, anything might match this position
        _ -> RML_WildCard

typeToRoughMatchTc :: Type -> RoughMatchTc
typeToRoughMatchTc ty
  | Just (ty', _) <- splitCastTy_maybe ty   = typeToRoughMatchTc ty'
  | Just (tc,_)   <- splitTyConApp_maybe ty
  , not (isTypeFamilyTyCon tc)              = RM_KnownTc $! roughMatchTyConName tc
    -- See Note [Rough matching in class and family instances]
  | otherwise                               = RM_WildCard

roughMatchTyConName :: TyCon -> Name
roughMatchTyConName tc
  | tc_name == cONSTRAINTTyConName
  = tYPETyConName  -- TYPE and CONSTRAINT are not apart, so they must use
                   -- the same rough-map key. We arbitrarily use TYPE.
                   -- See Note [Type and Constraint are not apart]
                   -- wrinkle (W1) in GHC.Builtin.Types.Prim
  | otherwise
  = assertPpr (isGenerativeTyCon tc Nominal) (ppr tc) tc_name
  where
    tc_name = tyConName tc


-- | Trie of @[RoughMatchTc]@
--
-- *Examples*
-- @
-- insert [OtherTc] 1
-- insert [OtherTc] 2
-- lookup [OtherTc] == [1,2]
-- @
data RoughMap a
  = RMEmpty -- An optimised (finite) form of emptyRM
            -- Invariant: Empty RoughMaps are always represented with RMEmpty

  | RM { rm_empty :: Bag a
           -- Keyed by an empty [RoughMapTc]

       , rm_known :: DNameEnv (RoughMap a)
           -- Keyed by (RM_KnownTc tc : rm_tcs)
           -- DNameEnv: see Note [InstEnv determinism] in GHC.Core.InstEnv

       , rm_wild :: RoughMap a }
           -- Keyed by (RM_WildCard : rm_tcs)
  deriving (Functor)

instance Outputable a => Outputable (RoughMap a) where
  ppr (RM empty known unknown) =
      vcat [text "RM"
           , nest 2 (vcat [ text "Empty:" <+> ppr empty
                          , text "Known:" <+> ppr known
                          , text "Unknown:" <+> ppr unknown])]
  ppr RMEmpty = text "{}"

emptyRM :: RoughMap a
emptyRM = RMEmpty

-- | Order of result is deterministic.
lookupRM :: [RoughMatchLookupTc] -> RoughMap a -> [a]
lookupRM tcs rm = bagToList (fst $ lookupRM' tcs rm)


-- | N.B. Returns a 'Bag' for matches, which allows us to avoid rebuilding all of the lists
-- we find in 'rm_empty', which would otherwise be necessary due to '++' if we
-- returned a list. We use a list for unifiers because the tail is computed lazily and
-- we often only care about the first couple of potential unifiers. Constructing a
-- bag forces the tail which performs much too much work.
--
-- See Note [Matching a RoughMap]
-- See Note [Matches vs Unifiers]
lookupRM' :: [RoughMatchLookupTc] -> RoughMap a -> (Bag a -- Potential matches
                                                   , [a]) -- Potential unifiers
lookupRM' _ RMEmpty   -- The RoughMap is empty
  = (emptyBag, [])

lookupRM' [] rm       -- See Note [Simple Matching Semantics] about why
  = (listToBag m, m)  -- we return everything when the lookup key runs out
  where
    m = elemsRM rm

lookupRM' (RML_KnownTc tc : tcs) rm  =
  let (common_m, common_u) = lookupRM' tcs (rm_wild rm)
      (m, u) = maybe (emptyBag, []) (lookupRM' tcs) (lookupDNameEnv (rm_known rm) tc)
  in ( rm_empty rm `unionBags` common_m `unionBags` m
     , bagToList (rm_empty rm) ++ common_u ++ u)

-- A RML_NoKnownTC does **not** match any KnownTC but can unify
lookupRM' (RML_NoKnownTc : tcs) rm =
  let (u_m, _u_u) = lookupRM' tcs (rm_wild rm)
  in ( rm_empty rm `unionBags` u_m -- Definitely don't match
     , snd $ lookupRM' (RML_WildCard : tcs) rm) -- But could unify..

lookupRM' (RML_WildCard : tcs)    rm  =
--  pprTrace "RM wild" (ppr tcs $$ ppr (eltsDNameEnv (rm_known rm))) $
  let (m, u)     = foldDNameEnv add_one (emptyBag, []) (rm_known rm)
      (u_m, u_u) = lookupRM' tcs (rm_wild rm)
  in ( rm_empty rm `unionBags` u_m `unionBags` m
     , bagToList (rm_empty rm) ++ u_u ++ u )
  where
     add_one :: RoughMap a -> (Bag a, [a]) -> (Bag a, [a])
     add_one rm ~(m2, u2) = (m1 `unionBags` m2, u1 ++ u2)
                          where
                            (m1,u1) = lookupRM' tcs rm

unionRM :: RoughMap a -> RoughMap a -> RoughMap a
unionRM RMEmpty a = a
unionRM a RMEmpty = a
unionRM a b =
  RM { rm_empty = rm_empty a `unionBags` rm_empty b
     , rm_known = plusDNameEnv_C unionRM (rm_known a) (rm_known b)
     , rm_wild = rm_wild a `unionRM` rm_wild b
     }


insertRM :: [RoughMatchTc] -> a -> RoughMap a -> RoughMap a
insertRM k v RMEmpty =
    insertRM k v $ RM { rm_empty = emptyBag
                      , rm_known = emptyDNameEnv
                      , rm_wild = emptyRM }
insertRM [] v rm@(RM {}) =
    -- See Note [Simple Matching Semantics]
    rm { rm_empty = v `consBag` rm_empty rm }

insertRM (RM_KnownTc k : ks) v rm@(RM {}) =
    rm { rm_known = alterDNameEnv f (rm_known rm) k }
  where
    f Nothing  = Just $ (insertRM ks v emptyRM)
    f (Just m) = Just $ (insertRM ks v m)

insertRM (RM_WildCard : ks) v rm@(RM {}) =
    rm { rm_wild = insertRM ks v (rm_wild rm) }

filterRM :: (a -> Bool) -> RoughMap a -> RoughMap a
filterRM _ RMEmpty = RMEmpty
filterRM pred rm =
    normalise $ RM {
      rm_empty = filterBag pred (rm_empty rm),
      rm_known = mapDNameEnv (filterRM pred) (rm_known rm),
      rm_wild = filterRM pred (rm_wild rm)
    }

-- | Place a 'RoughMap' in normal form, turning all empty 'RM's into
-- 'RMEmpty's. Necessary after removing items.
normalise :: RoughMap a -> RoughMap a
normalise RMEmpty = RMEmpty
normalise (RM empty known RMEmpty)
  | isEmptyBag empty
  , isEmptyDNameEnv known = RMEmpty
normalise rm = rm

-- | Filter all elements that might match a particular key with the given
-- predicate.
filterMatchingRM :: (a -> Bool) -> [RoughMatchTc] -> RoughMap a -> RoughMap a
filterMatchingRM _    _  RMEmpty = RMEmpty
filterMatchingRM pred [] rm      = filterRM pred rm
filterMatchingRM pred (RM_KnownTc tc : tcs) rm =
    normalise $ RM {
      rm_empty = filterBag pred (rm_empty rm),
      rm_known = alterDNameEnv (join . fmap (dropEmpty . filterMatchingRM pred tcs)) (rm_known rm) tc,
      rm_wild = filterMatchingRM pred tcs (rm_wild rm)
    }
filterMatchingRM pred (RM_WildCard : tcs) rm =
    normalise $ RM {
      rm_empty = filterBag pred (rm_empty rm),
      rm_known = mapDNameEnv (filterMatchingRM pred tcs) (rm_known rm),
      rm_wild = filterMatchingRM pred tcs (rm_wild rm)
    }

dropEmpty :: RoughMap a -> Maybe (RoughMap a)
dropEmpty RMEmpty = Nothing
dropEmpty rm = Just rm

elemsRM :: RoughMap a -> [a]
elemsRM = foldRM (:) []

foldRM :: (a -> b -> b) -> b -> RoughMap a -> b
foldRM f = go
  where
    -- N.B. local worker ensures that the loop can be specialised to the fold
    -- function.
    go z RMEmpty = z
    go z (RM{ rm_wild = unk, rm_known = known, rm_empty = empty}) =
      foldr
        f
        (foldDNameEnv
           (flip go)
           (go z unk)
           known
        )
        empty

nonDetStrictFoldRM :: (b -> a -> b) -> b -> RoughMap a -> b
nonDetStrictFoldRM f = go
  where
    -- N.B. local worker ensures that the loop can be specialised to the fold
    -- function.
    go !z RMEmpty = z
    go  z rm@(RM{}) =
      foldl'
        f
        (nonDetStrictFoldDNameEnv
           (flip go)
           (go z (rm_wild rm))
           (rm_known rm)
        )
        (rm_empty rm)

sizeRM :: RoughMap a -> Int
sizeRM = nonDetStrictFoldRM (\acc _ -> acc + 1) 0