CodeGen.hs

{-# LANGUAGE FlexibleContexts  #-}
{-# LANGUAGE OverloadedLabels  #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TupleSections     #-}
{-# LANGUAGE TypeApplications  #-}

module Wingman.CodeGen
  ( module Wingman.CodeGen
  , module Wingman.CodeGen.Utils
  ) where


import           ConLike
import           Control.Lens ((%~), (<>~), (&))
import           Control.Monad.Except
import           Control.Monad.Reader (ask)
import           Control.Monad.State
import           Data.Bool (bool)
import           Data.Functor ((<&>))
import           Data.Generics.Labels ()
import           Data.List
import qualified Data.Set as S
import           Data.Traversable
import           DataCon
import           Development.IDE.GHC.Compat
import           GHC.Exts
import           GHC.SourceGen (occNameToStr)
import           GHC.SourceGen.Binds
import           GHC.SourceGen.Expr
import           GHC.SourceGen.Overloaded
import           GHC.SourceGen.Pat
import           GhcPlugins (isSymOcc, mkVarOccFS)
import           OccName (occName)
import           PatSyn
import           Type hiding (Var)
import           Wingman.CodeGen.Utils
import           Wingman.GHC
import           Wingman.Judgements
import           Wingman.Judgements.Theta
import           Wingman.Machinery
import           Wingman.Naming
import           Wingman.Types


destructMatches
    :: Bool
    -> (ConLike -> Judgement -> Rule)
       -- ^ How to construct each match
    -> Maybe OccName
       -- ^ Scrutinee
    -> CType
       -- ^ Type being destructed
    -> Judgement
    -> RuleM (Synthesized [RawMatch])
-- TODO(sandy): In an ideal world, this would be the same codepath as
-- 'destructionFor'. Make sure to change that if you ever change this.
destructMatches use_field_puns f scrut t jdg = do
  let hy = jEntireHypothesis jdg
      g  = jGoal jdg
  case tacticsGetDataCons $ unCType t of
    Nothing -> throwError $ GoalMismatch "destruct" g
    Just (dcs, apps) ->
      fmap unzipTrace $ for dcs $ \dc -> do
        let con = RealDataCon dc
            ev = concatMap mkEvidence $ dataConInstArgTys dc apps
            -- We explicitly do not need to add the method hypothesis to
            -- #syn_scoped
            method_hy = foldMap evidenceToHypothesis ev
            args = conLikeInstOrigArgTys' con apps
        modify $ evidenceToSubst ev
        subst <- gets ts_unifier
        ctx <- ask

        let names_in_scope = hyNamesInScope hy
            names = mkManyGoodNames (hyNamesInScope hy) args
            (names', destructed) =
              mkDestructPat (bool Nothing (Just names_in_scope) use_field_puns) con names

        let hy' = patternHypothesis scrut con jdg
                $ zip names'
                $ coerce args
            j = fmap (CType . substTyAddInScope subst . unCType)
              $ introduce ctx hy'
              $ introduce ctx method_hy
              $ withNewGoal g jdg
        ext <- f con j
        pure $ ext
          & #syn_trace %~ rose ("match " <> show dc <> " {" <> intercalate ", " (fmap show names') <> "}")
                        . pure
          & #syn_scoped <>~ hy'
          & #syn_val %~ match [destructed] . unLoc


------------------------------------------------------------------------------
-- | Generate just the 'Match'es for a case split on a specific type.
destructionFor :: Hypothesis a -> Type -> Maybe [LMatch GhcPs (LHsExpr GhcPs)]
-- TODO(sandy): In an ideal world, this would be the same codepath as
-- 'destructMatches'. Make sure to change that if you ever change this.
destructionFor hy t = do
  case tacticsGetDataCons t of
    Nothing -> Nothing
    Just ([], _) -> Nothing
    Just (dcs, apps) -> do
      for dcs $ \dc -> do
        let con   = RealDataCon dc
            args  = conLikeInstOrigArgTys' con apps
            names = mkManyGoodNames (hyNamesInScope hy) args
        pure
          . noLoc
          . Match
              noExtField
              CaseAlt
              [toPatCompat $ snd $ mkDestructPat Nothing con names]
          . GRHSs noExtField (pure $ noLoc $ GRHS noExtField [] $ noLoc $ var "_")
          . noLoc
          $ EmptyLocalBinds noExtField



------------------------------------------------------------------------------
-- | Produces a pattern for a data con and the names of its fields.
mkDestructPat :: Maybe (S.Set OccName) -> ConLike -> [OccName] -> ([OccName], Pat GhcPs)
mkDestructPat already_in_scope con names
  | RealDataCon dcon <- con
  , isTupleDataCon dcon =
      (names, tuple pat_args)
  | fields@(_:_) <- zip (conLikeFieldLabels con) names
  , Just in_scope <- already_in_scope =
      let (names', rec_fields) =
            unzip $ fields <&> \(label, name) -> do
              let label_occ = mkVarOccFS $ flLabel label
              case S.member label_occ in_scope of
                -- We have a shadow, so use the generated name instead
                True ->
                  (name,) $ noLoc $
                    HsRecField
                      (noLoc $ mkFieldOcc $ noLoc $ Unqual label_occ)
                      (noLoc $ bvar' name)
                      False
                -- No shadow, safe to use a pun
                False ->
                  (label_occ,) $ noLoc $
                    HsRecField
                      (noLoc $ mkFieldOcc $ noLoc $ Unqual label_occ)
                      (noLoc $ bvar' label_occ)
                      True

        in (names', )
         $ ConPatIn (noLoc $ Unqual $ occName $ conLikeName con)
         $ RecCon
         $ HsRecFields rec_fields
         $ Nothing
  | otherwise =
      (names, ) $ infixifyPatIfNecessary con $
        conP
          (coerceName $ conLikeName con)
          pat_args
  where
    pat_args = fmap bvar' names


infixifyPatIfNecessary :: ConLike -> Pat GhcPs -> Pat GhcPs
infixifyPatIfNecessary dcon x
  | conLikeIsInfix dcon =
      case x of
        ConPatIn op (PrefixCon [lhs, rhs]) ->
          ConPatIn op $ InfixCon lhs rhs
        y -> y
  | otherwise = x



unzipTrace :: [Synthesized a] -> Synthesized [a]
unzipTrace = sequenceA


-- | Essentially same as 'dataConInstOrigArgTys' in GHC,
--  but only accepts universally quantified types as the second arguments
--  and automatically introduces existentials.
--
-- NOTE: The behaviour depends on GHC's 'dataConInstOrigArgTys'.
--       We need some tweaks if the compiler changes the implementation.
conLikeInstOrigArgTys'
  :: ConLike
      -- ^ 'DataCon'structor
  -> [Type]
      -- ^ /Universally/ quantified type arguments to a result type.
      --   It /MUST NOT/ contain any dictionaries, coercion and existentials.
      --
      --   For example, for @MkMyGADT :: b -> MyGADT a c@, we
      --   must pass @[a, c]@ as this argument but not @b@, as @b@ is an existential.
  -> [Type]
      -- ^ Types of arguments to the ConLike with returned type is instantiated with the second argument.
conLikeInstOrigArgTys' con uniTys =
  let exvars = conLikeExTys con
   in conLikeInstOrigArgTys con $
        uniTys ++ fmap mkTyVarTy exvars
      -- Rationale: At least in GHC <= 8.10, 'dataConInstOrigArgTys'
      -- unifies the second argument with DataCon's universals followed by existentials.
      -- If the definition of 'dataConInstOrigArgTys' changes,
      -- this place must be changed accordingly.


conLikeExTys :: ConLike -> [TyCoVar]
conLikeExTys (RealDataCon d) = dataConExTys d
conLikeExTys (PatSynCon p) = patSynExTys p

patSynExTys :: PatSyn -> [TyCoVar]
patSynExTys ps = patSynExTyVars ps


------------------------------------------------------------------------------
-- | Combinator for performing case splitting, and running sub-rules on the
-- resulting matches.

destruct' :: Bool -> (ConLike -> Judgement -> Rule) -> HyInfo CType -> Judgement -> Rule
destruct' use_field_puns f hi jdg = do
  when (isDestructBlacklisted jdg) $ throwError NoApplicableTactic
  let term = hi_name hi
  ext
      <- destructMatches
           use_field_puns
           f
           (Just term)
           (hi_type hi)
           $ disallowing AlreadyDestructed (S.singleton term) jdg
  pure $ ext
    & #syn_trace     %~ rose ("destruct " <> show term) . pure
    & #syn_used_vals %~ S.insert term
    & #syn_val       %~ noLoc . case' (var' term)


------------------------------------------------------------------------------
-- | Combinator for performign case splitting, and running sub-rules on the
-- resulting matches.
destructLambdaCase' :: Bool -> (ConLike -> Judgement -> Rule) -> Judgement -> Rule
destructLambdaCase' use_field_puns f jdg = do
  when (isDestructBlacklisted jdg) $ throwError NoApplicableTactic
  let g  = jGoal jdg
  case splitFunTy_maybe (unCType g) of
    Just (arg, _) | isAlgType arg ->
      fmap (fmap noLoc lambdaCase) <$>
        destructMatches use_field_puns f Nothing (CType arg) jdg
    _ -> throwError $ GoalMismatch "destructLambdaCase'" g


------------------------------------------------------------------------------
-- | Construct a data con with subgoals for each field.
buildDataCon
    :: Bool       -- Should we blacklist destruct?
    -> Judgement
    -> ConLike            -- ^ The data con to build
    -> [Type]             -- ^ Type arguments for the data con
    -> RuleM (Synthesized (LHsExpr GhcPs))
buildDataCon should_blacklist jdg dc tyapps = do
  args <- case dc of
    RealDataCon dc' -> do
      let (skolems', theta, args) = dataConInstSig dc' tyapps
      modify $ \ts ->
        evidenceToSubst (foldMap mkEvidence theta) ts
          & #ts_skolems <>~ S.fromList skolems'
      pure args
    _ ->
      -- If we have a 'PatSyn', we can't continue, since there is no
      -- 'dataConInstSig' equivalent for 'PatSyn's. I don't think this is
      -- a fundamental problem, but I don't know enough about the GHC internals
      -- to implement it myself.
      --
      -- Fortunately, this isn't an issue in practice, since 'PatSyn's are
      -- never in the hypothesis.
      throwError $ TacticPanic "Can't build Pattern constructors yet"
  ext
      <- fmap unzipTrace
       $ traverse ( \(arg, n) ->
                    newSubgoal
                  . filterSameTypeFromOtherPositions dc n
                  . bool id blacklistingDestruct should_blacklist
                  . flip withNewGoal jdg
                  $ CType arg
                  ) $ zip args [0..]
  pure $ ext
    & #syn_trace %~ rose (show dc) . pure
    & #syn_val   %~ mkCon dc tyapps


------------------------------------------------------------------------------
-- | Make a function application, correctly handling the infix case.
mkApply :: OccName -> [HsExpr GhcPs] -> LHsExpr GhcPs
mkApply occ (lhs : rhs : more)
  | isSymOcc occ
  = noLoc $ foldl' (@@) (op lhs (coerceName occ) rhs) more
mkApply occ args = noLoc $ foldl' (@@) (var' occ) args


------------------------------------------------------------------------------
-- | Run a tactic over each term in the given 'Hypothesis', binding the results
-- of each in a let expression.
letForEach
    :: (OccName -> OccName)           -- ^ How to name bound variables
    -> (HyInfo CType -> TacticsM ())  -- ^ The tactic to run
    -> Hypothesis CType               -- ^ Terms to generate bindings for
    -> Judgement                      -- ^ The goal of original hole
    -> RuleM (Synthesized (LHsExpr GhcPs))
letForEach rename solve (unHypothesis -> hy) jdg = do
  case hy of
    [] -> newSubgoal jdg
    _ -> do
      ctx <- ask
      let g = jGoal jdg
      terms <- fmap sequenceA $ for hy $ \hi -> do
        let name = rename $ hi_name hi
        res <- tacticToRule jdg $ solve hi
        pure $ fmap ((name,) . unLoc) res
      let hy' = fmap (g <$) $ syn_val terms
          matches = fmap (fmap (\(occ, expr) -> valBind (occNameToStr occ) expr)) terms
      g <- fmap (fmap unLoc) $ newSubgoal $ introduce ctx (userHypothesis hy') jdg
      pure $ fmap noLoc $ let' <$> matches <*> g