TyConSubst.hs

{-# LANGUAGE CPP #-}

{-# LANGUAGE LambdaCase          #-}
{-# LANGUAGE NamedFieldPuns      #-}
{-# LANGUAGE RecordWildCards     #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TupleSections       #-}
{-# LANGUAGE ViewPatterns        #-}

{-|
Module: GHC.TcPlugin.API.TyConSubst

This module provides functionality for recognising whether a type is a
'TyConApp' while taking into account Given nominal equalities.

In particular, this allows dealing with flattening variables on older GHC versions
(9.0 and below). For example, instead of @[W] m + n ~ n + m@, older GHCs might
produce @[G] u ~ m + n, [G] v ~ n + m, [W] u ~ v@. In this case, one needs to
use the Givens to recognise that @u ~ v@ can be solved using the laws of natural
number arithmetic.

Usage:

  - Use 'mkTyConSubst' to create a 'TyConSubst' from Givens.
  - Use 'splitTyConApp_upTo' to compute whether a type is a 'TyConApp', taking
    into account Given constraints (in the form of the 'TyConSubst').

Notes:

  - 'splitTyConApp_upTo' will also look through type synonyms,
  - 'splitTyConApp_upTo' only takes into account nominal equalities.
    Please open a ticket if you have a need for rewriting modulo representational
    equalities.

-}
module GHC.TcPlugin.API.TyConSubst (
    TyConSubst -- opaque
  , mkTyConSubst
  , splitTyConApp_upTo
  ) where

-- base
import Data.Bifunctor
import Data.Either
  ( partitionEithers )
import Data.Foldable
  ( toList, asum )
import Data.List.NonEmpty
  ( NonEmpty(..) )

-- array
import qualified Data.Array as Array
  ( (!) )

-- containers
import Data.Graph
  ( Graph, Vertex )
import qualified Data.Graph as Graph
import Data.Map
  ( Map )
import qualified Data.Map as Map
import Data.Set
  ( Set )
import qualified Data.Set as Set

-- ghc
import GHC.Utils.Outputable
  hiding ( (<>) )

-- ghc-tcplugin-api
import GHC.TcPlugin.API
import GHC.Tc.Types.Constraint

{-------------------------------------------------------------------------------
  The main type

  TODO: maybe this could be sped up with
  <https://hackage.haskell.org/package/union-find>?
-------------------------------------------------------------------------------}

-- | Substitution for recognizing 'TyCon' applications modulo nominal equalities.
data TyConSubst = TyConSubst {
      tyConSubstMap :: Map TcTyVar (NonEmpty (TyCon, [Type], [Coercion]))
    , tyConSubstCanon :: Map TcTyVar (TcTyVar, [Coercion])
    }
-- During constraint solving the set of Given constraints includes so-called
-- "canonical equalities": equalities of the form
--
-- > var ~ typ                  (CTyEqCan)
-- > var ~ TyCon arg1 .. argN   (CFunEqCan, the TyCon will be a type family)
--
-- The problem we want to solve is recognizing if some type τ is of the form
--
-- > TyCon arg1 arg2 .. argN   (0 <= N)
--
-- modulo those canonical equalities. We limit the scope of what we try to do:
--
-- o We are only interested in recognizing types of the form above
--   (as opposed to general parsing-modulo-equalities).
-- o We will only use the canonical equalities as-is: we will not attempt to
--   derive any additional equalities from them (i.e. if, say, we know that
--   @x ~ T1@ and @x ~ T2@, we will not attempt to use the fact that this means
--   that @T1 ~ T2@, nor any derived conclusions thereof). We /will/ however
--   try to apply the canononical equalities as often as is necessary (e.g.,
--   first applying @x ~ T y@, then applying @y ~ T2@).
--
-- We solve this problem by constructing a 'TyConSubst': a possibly
-- non-deterministic substitution mapping type variables to types of the form
-- above (that is, a type constructor applied to some arguments).
--
-- We detail the construction of this substitution below (see documentation of
-- 'Classified' and 'process'), but once we have this substitution, the
-- recognition problem becomes easy:
--
-- 1. Without loss of generality, let τ be of the form @t arg1 arg2 .. argN@
-- 2. If @t@ is a 'TyCon', we're done.
-- 3. Otherwise, if @t@ is a variable @x@, lookup @x@ in the substitution; if
--    there is one (or more) mappings for @x@, then we have successfully
--    recognized τ to be of the form above. There is no need to apply the
--    substitution repeatedly.
--
-- The substitution is non-deterministic because there might be multiple
-- matches. For example, if we have
--
-- > type family Foo where
-- >   Foo = Int
--
-- then we might well have equalities @x ~ Int, x ~ Foo@ in scope, and so a type
-- @x@ would match two different 'TyCon's. What we do know, however, is that if
-- τ matches both @t arg1 .. argN@ and @t' arg1' .. argM'@ (possibly @N /= M@),
-- then
--
-- > t arg1 .. argN ~ t' arg1' .. argM'
--
-- If @t == t'@, we can conclude that the arguments are equal only if @t@ is
-- injective.

{-------------------------------------------------------------------------------
  Basic functionality for working with 'TyConSubst'
-------------------------------------------------------------------------------}

-- | Empty substitution
--
-- The canonical variables map is established once when the initial substitution
-- is generated and not updated thereafter.
tyConSubstEmpty :: Map TcTyVar (TcTyVar, [Coercion]) -> TyConSubst
tyConSubstEmpty canon = TyConSubst {
      tyConSubstMap   = Map.empty
    , tyConSubstCanon = canon
    }

-- | Lookup a variable in the substitution
tyConSubstLookup :: TcTyVar -> TyConSubst -> Maybe (NonEmpty (TyCon, [Type], [Coercion]))
tyConSubstLookup var TyConSubst{..} =
  fmap ( \ (tc, tys, deps) -> (tc, tys, deps ++ deps')) <$> Map.lookup var' tyConSubstMap
  where
    var' :: TcTyVar
    deps' :: [Coercion]
    (var', deps') = canonicalize tyConSubstCanon var

-- | Extend substitution with new bindings
tyConSubstExtend ::
     [(TcTyVar, (TyCon, [Type]), [Coercion])]
  -> TyConSubst -> TyConSubst
tyConSubstExtend new subst@TyConSubst{..} = subst {
      tyConSubstMap = Map.unionWith (<>)
                        (Map.fromList $ map aux new)
                        tyConSubstMap
    }
  where
    aux :: (TcTyVar,          (TyCon, [Type]), [Coercion])
        -> (TcTyVar, NonEmpty (TyCon, [Type] , [Coercion]))
    aux (var, (tc, args), deps) =
      let (var', deps') = canonicalize tyConSubstCanon var
      in  (var', (tc, args, deps ++ deps') :| [])

{-------------------------------------------------------------------------------
  Classification
-------------------------------------------------------------------------------}

-- | Classified canonical nominal equality constraints.
--
-- The first step in the construction of the 'TyConSubst' is to classify the
-- available canonical equalities as one of three categories, defined below.
data Classified = Classified {
      -- | " Obviously " productive mappings
      --
      -- An equality @var := TyCon args@ is productive, because as soon as we
      -- apply it, we are done: we have successfully recognized a type as being
      -- an application of a concrete type constructor (note that we only ever
      -- apply the substitution to the head @t@ of a type @t args@, never to the
      -- arguments).
      classifiedProductive :: [(TcTyVar, (TyCon, [Type]), [Coercion])]

      -- | Extend equivalence class of variables
      --
      -- An equality @var1 := var2@ we will regard as extending the equivalence
      -- classes of variables (see 'constructEquivClasses').
    , classifiedExtendEquivClass :: [(TcTyVar, TcTyVar, [Coercion])]

      -- | Substitutions we need to reconsider later
      --
      -- An equality @var1 := var2 args@ (with @args@ a non-empty list of
      -- arguments) is most problematic. Applying it /may/ allow us to make
      -- progress, but it may not (consider for example @var := var arg@). We
      -- will reconsider such equalities at the end (see 'process').
    , classifiedReconsider :: [(TcTyVar, (TcTyVar, NonEmpty Type), [Coercion])]
    }

instance Semigroup Classified where
  c1 <> c2 = Classified {
        classifiedProductive       = combine classifiedProductive
      , classifiedExtendEquivClass = combine classifiedExtendEquivClass
      , classifiedReconsider       = combine classifiedReconsider
      }
    where
      combine :: (Classified -> [a]) -> [a]
      combine f = f c1 ++ f c2

instance Monoid Classified where
  mempty = Classified [] [] []

productive :: TcTyVar -> (TyCon, [Type]) -> [Coercion] -> Classified
productive var app deps = mempty {
      classifiedProductive = [(var, app, deps)]
    }

extendEquivClass :: TcTyVar -> TcTyVar -> [Coercion] -> Classified
extendEquivClass var var' deps = mempty {
      classifiedExtendEquivClass = [(var, var', deps)]
    }

reconsider :: TcTyVar -> (TcTyVar, NonEmpty Type) -> [Coercion] -> Classified
reconsider var (var', args) deps = mempty {
      classifiedReconsider = [(var, (var', args), deps)]
    }

-- | Classify a set of Given constraints.
--
-- See 'Classified' for details.
classify :: [Ct] -> Classified
classify = go mempty
  where
    go :: Classified -> [Ct] -> Classified
    go acc []     = acc
    go acc (c:cs) =
      let deps = [ctEvCoercion (ctEvidence c)] in
        case isCanonicalVarEq c of
          Just (var, splitAppTys -> (fn, args), NomEq)
            | Just (tyCon, inner) <- splitTyConApp_maybe fn ->
                go (productive var (tyCon, inner ++ args) deps <> acc) cs
            | Just var' <- getTyVar_maybe fn, null args ->
                go (extendEquivClass var var' deps <> acc) cs
            | Just var' <- getTyVar_maybe fn, x:xs <- args ->
                go (reconsider var (var', x :| xs) deps <> acc) cs
          _otherwise ->
            go acc cs

{-------------------------------------------------------------------------------
  Processing
-------------------------------------------------------------------------------}

-- | Construct 'TyCon' substitution from classified nominal equality constraints.
--
-- The difficult part in constructing this substitution are the equalities of
-- the form @var1 ~ var2 args@, which we ear-marked as "to reconsider" during
-- classification.
--
-- We will do this iteratively:
--
-- o We first construct a set of variable equivalence classes based on
--   'classifiedExtendEquivClass' (using 'constructEquivClasses'), and use that
--   along with the "obviously productive" equalities ('classifiedProductive')
--   as the initial value of the accumulator (a 'TyConSubst').
-- o We then repeatedly consider the remaining equalities. Whenever there is
--   a substitution available in the accumulator for @var2@ which turns it into
--   a type of the form @TyCon args'@, we add @var1 := TyCon args' args@ to the
--   accumulator.
-- o We keep doing this until we can make no more progress.
--
-- The functions for working with 'TyConSubst' take the variable equivalence
-- classes into acocunt, so we do not need to do that here.
--
-- Two observations:
--
-- o This process must terminate: there are a finite number of constraints
--   to consider, and whenever we apply a substitution from the accumulator,
--   we get an "obviously productive" substitution: we do not create new work
--   in the loop.
-- o We may end up ignoring some substitutions: if there is a substitution
--   @var1 := var2 args@ and we don't have any (productive) substitutions for
--   @var2@, we will just ignore it.
--
-- A note on recursive bindings: a direct or indirect recursive binding
--
-- > x := x args1      x := y args1
-- >                   y := x args2
--
-- where @args1, args2@ are non-empty lists of arguments, /cannot/ be relevant:
-- if they were, that would imply that there is some type constructor (regular
-- datatype or type family) which can be applied to an arbitrary number of
-- arguments. Such datatypes or type families cannot be defined in Haskell.
-- We therefore take no special care in handling recursive bindings, other than
-- to note (as we did above) that the process must terminate.
process :: Classified -> TyConSubst
process Classified{..} =
    go initSubst classifiedReconsider
  where
    initSubst :: TyConSubst
    initSubst =
          tyConSubstExtend classifiedProductive
        $ tyConSubstEmpty (constructEquivClasses classifiedExtendEquivClass)

    go :: TyConSubst
       -> [(TcTyVar, (TcTyVar, NonEmpty Type), [Coercion])]
       -> TyConSubst
    go acc rs =
        let (prod, rest) = tryApply makeProductive rs in
        if null prod
          then acc -- No other equations can be made productive
          else go (tyConSubstExtend prod acc) rest
      where
        makeProductive ::
             (TcTyVar, (TcTyVar, NonEmpty Type), [Coercion])
          -> Maybe (NonEmpty (TcTyVar, (TyCon, [Type]), [Coercion]))
        makeProductive (var, (var', args), deps) = do
          tcApp <- tyConSubstLookup var' acc
          return $ fmap aux tcApp
          where
            aux :: (TyCon, [Type], [Coercion]) -> (TcTyVar, (TyCon, [Type]), [Coercion])
            aux (tyCon, args', deps') =
              (var, (tyCon, args' ++ toList args), deps ++ deps')

-- | Construct a 'TyConSubst' from a collection of Given constraints.
mkTyConSubst :: [Ct] -> TyConSubst
mkTyConSubst = process . classify

{-------------------------------------------------------------------------------
  Using
-------------------------------------------------------------------------------}

-- | Like 'splitTyConApp_maybe', but taking Given constraints into account.
--
-- Alongside the @TyCon@ and its arguments, also returns a list of coercions
-- that embody the Givens that we depended on.
--
-- Looks through type synonyms, just like 'splitTyConApp_maybe' does.
splitTyConApp_upTo :: TyConSubst -> Type -> Maybe (NonEmpty (TyCon, [Type], [Coercion]))
splitTyConApp_upTo subst typ = asum [
      -- Direct match
      do (tyCon, inner) <- splitTyConApp_maybe fn
         return ((tyCon, inner ++ args, []) :| [])

      -- Indirect match
    , do var <- getTyVar_maybe fn
         tcApps <- tyConSubstLookup var subst
         return $
           fmap (\ (tc, inner, deps) -> (tc, inner ++ args, deps)) tcApps
    ]
  where
    (fn, args) = splitAppTys typ

{-------------------------------------------------------------------------------
  Outputable
-------------------------------------------------------------------------------}

instance Outputable TyConSubst where
  ppr TyConSubst{..} = parens $
          text "TyConSubst"
      <+> ppr tyConSubstMap
      <+> ppr tyConSubstCanon

{-------------------------------------------------------------------------------
  Canonical equalities
-------------------------------------------------------------------------------}

isCanonicalVarEq :: Ct -> Maybe (TcTyVar, Type, EqRel)
isCanonicalVarEq = \case
#if __GLASGOW_HASKELL__ < 902
    CTyEqCan { cc_tyvar, cc_rhs, cc_eq_rel } ->
      Just (cc_tyvar, cc_rhs, cc_eq_rel)
    CFunEqCan { cc_fsk, cc_fun, cc_tyargs } ->
      Just (cc_fsk, mkTyConApp cc_fun cc_tyargs, NomEq)
    _otherwise    -> Nothing
#elif __GLASGOW_HASKELL__ < 907
    CEqCan { cc_lhs, cc_rhs, cc_eq_rel }
      | TyVarLHS var <- cc_lhs
      -> Just (var, cc_rhs, cc_eq_rel)
      | TyFamLHS tyCon args <- cc_lhs
      , Just var            <- getTyVar_maybe cc_rhs
      -> Just (var, mkTyConApp tyCon args, NomEq)
    _otherwise
      -> Nothing
#else
    CEqCan eqCt
      | TyVarLHS var <- lhs
      -> Just (var, rhs, rel)
      | TyFamLHS tyCon args <- lhs
      , Just var            <- getTyVar_maybe rhs
      -> Just (var, mkTyConApp tyCon args, rel)
      where
        lhs = eq_lhs eqCt
        rhs = eq_rhs eqCt
        rel = eq_eq_rel eqCt
    _otherwise
      -> Nothing
#endif

{-------------------------------------------------------------------------------
  Internal auxiliary
-------------------------------------------------------------------------------}

-- | Attempt to apply a non-deterministic function to a list of values
--
-- Returns the successful results as well as the inputs on which the function
-- failed.
tryApply :: forall a b. (a -> Maybe (NonEmpty b)) -> [a] -> ([b], [a])
tryApply f = first (concat . map toList) . partitionEithers . map f'
  where
    f' :: a -> Either (NonEmpty b) a
    f' a = maybe (Right a) Left $ f a

{-------------------------------------------------------------------------------
  Equivalence classes
-------------------------------------------------------------------------------}

-- | Given a set of labelled equivalent pairs, map every value to a canonical
-- value, with the shortest path (as a list of labels) connecting the value to
-- the canonical value.
--
-- Example with two classes:
--
-- >>> constructEquivClasses [(1, 2, "12"), (4, 5, "45"), (2, 3, "23")]
-- fromList [(1,1,[]),(2,1,["12"]),(3,1, ["12","23"]),(4,4,[]),(5,4, ["45"])]
--
-- Adding one element that connects both classes:
--
-- >>> constructEquivClasses [(1, 2, "12"), (4, 5, "45"), (2, 3, "23"), (3,4,"34")]
-- fromList [(1,1,[]),(2,1,["12"]),(3,1,["12", "23"]),(4,1,["12","23","34"]),(5,4, ["12","23","34,"45"])]
constructEquivClasses :: forall a l. (Ord a, Monoid l) => [(a, a, l)] -> Map a (a, l)
constructEquivClasses equivs
  = Map.unions
  $ map (pickCanonical . map fromVertex . toList)
  $ Graph.components graph
  where
    allValues :: Set a
    allValues = Set.fromList $ concatMap (\(x, y, _) -> [x,y]) equivs

    edges :: Map (Set a) l
    edges = Map.fromList [ (Set.fromList [x, y], lbl) | (x,y,lbl) <- equivs ]

    toVertex   :: a -> Vertex
    fromVertex :: Vertex -> a

    toVertex   a = Map.findWithDefault (error "toVertex: impossible")   a $
                     Map.fromList $ zip (Set.elems allValues) [1..]
    fromVertex v = Map.findWithDefault (error "fromVertex: impossible") v $
                     Map.fromList $ zip [1..] (Set.elems allValues)

    graph :: Graph
    graph = Graph.buildG (1, Set.size allValues)
               [ (toVertex x, toVertex y) | (x, y, _) <- equivs]

    neighbours :: a -> [(a, l)]
    neighbours v =
      [ ( u, edges Map.! ( Set.fromList [ v, u ] ) )
      | u <- map fromVertex $ graph Array.! ( toVertex v )
      ]

    pickCanonical :: [a] -> Map a (a, l)
    pickCanonical comp = ( root, ) <$> go ( Map.singleton root mempty ) [ root ]
      where
        root = minimum comp

        go :: Map a l -> [a] -> Map a l
        go ds [] = ds
        go ds (v:vs) =
          let
            -- unvisited neighbours
            us = filter ( \ (u, _) -> not (u `Map.member` ds) ) $ neighbours v

            d = ds Map.! v
            ds' = Map.union ds (Map.fromList [ (u, l <> d) | (u, l) <- us ])
          in
            go ds' (vs ++ map fst us)

canonicalize :: (Ord a, Monoid l) => Map a (a, l) -> a -> (a, l)
canonicalize canon x = Map.findWithDefault (x, mempty) x canon