no-recursion-0.4.0.1: src/NoRecursion.hs

{-# LANGUAGE Unsafe #-}

-- |
-- Copyright: 2024 Greg Pfeil
-- License: AGPL-3.0-only WITH Universal-FOSS-exception-1.0 OR LicenseRef-proprietary
--
-- A plugin that identifies and reports on uses of recursion. The name evokes a
-- language pragma – implying a @Recursion@ pragma that is enabled by default.
module NoRecursion (plugin) where

import safe "base" Control.Applicative (liftA2, pure)
import safe "base" Control.Category ((.))
import safe "base" Control.Exception (ErrorCall (ErrorCall), throwIO)
import safe "base" Control.Monad ((=<<))
import safe "base" Data.Bool (Bool (False, True), not, (&&), (||))
import safe "base" Data.Either (Either (Left), either)
import safe "base" Data.Foldable
  ( all,
    any,
    elem,
    foldMap,
    foldr,
    foldrM,
    notElem,
    toList,
    traverse_,
  )
import safe "base" Data.Function (flip, ($))
import safe "base" Data.Functor (fmap, (<$>))
import safe "base" Data.Kind (Type)
import safe "base" Data.List (filter, intercalate, isPrefixOf, null)
import safe "base" Data.List.NonEmpty (NonEmpty, nonEmpty)
import safe "base" Data.Maybe (maybe)
import safe "base" Data.Semigroup (Semigroup ((<>)), (<>))
import safe "base" Data.String (String)
import safe "base" Data.Tuple (curry, fst, uncurry)
import "ghc" GHC.Plugins qualified as Plugins
import safe "this" PluginUtils
  ( Annotations,
    defaultPurePlugin,
    getAnnotations,
    processOptions,
  )

-- | The entrypoint for the "NoRecursion" plugin.
--
-- @since 0.1.0
plugin :: Plugins.Plugin
plugin =
  defaultPurePlugin
    { Plugins.installCoreToDos = \opts -> liftA2 install (parseOpts opts) . pure
    }

type Opts :: Type
data Opts = Opts
  { allowRecursion :: Bool,
    ignoreMethodCycles :: Bool,
    ignoredDecls :: [String],
    ignoredMethods :: [String]
  }

-- | The `Opts` we have if no @-fplugin-opts=NoRecursion:@ are provided.
--
-- - recursion is not allowed
-- - recursion cycles between methods is ignored (to avoid a breaking change)
defaultOpts :: Opts
defaultOpts =
  Opts
    { allowRecursion = False,
      ignoreMethodCycles = True,
      ignoredDecls = [],
      ignoredMethods = []
    }

type OptError :: Type
data OptError
  = MissingValue String
  | UnknownOption String
  | UnknownValue String String

prettyOptError :: OptError -> Plugins.SDoc
prettyOptError =
  Plugins.text . \case
    MissingValue opt ->
      "plugin option ‘NoRecursion:" <> opt <> "’ is missing a value"
    UnknownOption name -> "unknown plugin option ‘NoRecursion:" <> name <> "’"
    UnknownValue typ value ->
      "an option for the NoRecursion plugin was expecting a "
        <> typ
        <> " but received ‘"
        <> value
        <> "’"

parseBoolOpt :: String -> Either OptError Bool
parseBoolOpt = \case
  "true" -> pure True
  "false" -> pure False
  value -> Left $ UnknownValue "Bool" value

parseListOpt :: String -> [String]
parseListOpt =
  foldr
    ( curry $ \case
        (',', elems) -> [] : elems
        (c, []) -> [[c]]
        (c, curr : elems) -> (c : curr) : elems
    )
    []

parseOpts :: [Plugins.CommandLineOption] -> Plugins.CoreM Opts
parseOpts =
  foldrM
    ( \(name, mvalue) opts ->
        case name of
          "allow-recursion" ->
            either (err opts) (\v -> pure opts {allowRecursion = v}) $
              maybe (pure True) parseBoolOpt mvalue
          "ignore-method-cycles" ->
            either (err opts) (\v -> pure opts {ignoreMethodCycles = v}) $
              maybe (pure True) parseBoolOpt mvalue
          "ignore-decls" ->
            maybe
              (err opts $ MissingValue name)
              ( \v ->
                  pure opts {ignoredDecls = parseListOpt v <> ignoredDecls opts}
              )
              mvalue
          "ignore-methods" ->
            maybe
              (err opts $ MissingValue name)
              ( \v ->
                  pure
                    opts
                      { ignoredMethods = parseListOpt v <> ignoredMethods opts
                      }
              )
              mvalue
          _ -> err opts $ UnknownOption name
    )
    defaultOpts
    . processOptions
  where
    err opts = fmap (\() -> opts) . Plugins.errorMsg . prettyOptError

install :: Opts -> [Plugins.CoreToDo] -> [Plugins.CoreToDo]
install opts =
  (Plugins.CoreDoPluginPass "add NoRecursion rule" (noRecursionPass opts) :)

noRecursionPass :: Opts -> Plugins.ModGuts -> Plugins.CoreM Plugins.ModGuts
noRecursionPass opts guts = do
  dflags <- Plugins.getDynFlags
  anns <- getAnnotations guts
  either
    ( \recs ->
        Plugins.liftIO . throwIO . ErrorCall $
          "encountered recursion, which has been disabled:\n"
            <> intercalate "\n" (toList $ formatRecursionRecord dflags <$> recs)
    )
    (\() -> pure guts)
    . failOnRecursion dflags opts anns
    $ Plugins.mg_binds guts

type RecursionRecord :: Type -> Type
data RecursionRecord b = RecursionRecord [b] (NonEmpty b)

formatRecursionRecord ::
  (Plugins.Outputable b) => Plugins.DynFlags -> RecursionRecord b -> String
formatRecursionRecord dflags (RecursionRecord context recs) =
  maybe
    "at the top level"
    ( \v ->
        "in "
          <> intercalate
            " >> "
            (Plugins.showSDoc dflags . Plugins.ppr <$> toList v)
    )
    (nonEmpty context)
    <> ", the following bindings were recursive: "
    <> intercalate ", " (Plugins.showSDoc dflags . Plugins.ppr <$> toList recs)

recursionAnnotation :: String
recursionAnnotation = "Recursion"

noRecursionAnnotation :: String
noRecursionAnnotation = "NoRecursion"

moduleAllowsRecursion :: Bool -> [String] -> Bool
moduleAllowsRecursion allowRecursion modAnns =
  (allowRecursion || elem recursionAnnotation modAnns)
    && notElem noRecursionAnnotation modAnns

getName :: Plugins.DynFlags -> Plugins.CoreBndr -> String
getName dflags = Plugins.showSDoc dflags . Plugins.ppr

isInternalName :: Plugins.DynFlags -> Plugins.CoreBndr -> Bool
isInternalName dflags var =
  let v = getName dflags var
   in "$c" `isPrefixOf` v || "$f" `isPrefixOf` v

failOnRecursion ::
  Plugins.DynFlags ->
  Opts ->
  Annotations [String] ->
  [Plugins.CoreBind] ->
  Either (NonEmpty (RecursionRecord Plugins.CoreBndr)) ()
failOnRecursion
  dflags
  opts
  (modAnns, nameAnns)
  original =
    traverse_ Left
      . nonEmpty
      -- __TODO__: Default method implementations seem to cause mutual
      --           recursion with the instance, so here we filter them out,
      --           but this probably lets some real mutual recursion slip
      --           through.
      . filter
        ( not
            . \(RecursionRecord context recs) ->
              ignoreMethodCycles opts && null context && all (isInternalName dflags) recs
                || any (flip elem (ignoredDecls opts) . getName dflags) recs
                || any (flip elem (("$c" <>) <$> ignoredMethods opts) . getName dflags) context
        )
      $ recursiveCallsForBind
        =<< filter
          ( not
              . allowBind
                (moduleAllowsRecursion (allowRecursion opts) modAnns)
                nameAnns
          )
          original

addBindingReference :: b -> [RecursionRecord b] -> [RecursionRecord b]
addBindingReference var =
  fmap (\(RecursionRecord context recs) -> RecursionRecord (var : context) recs)

allowBind :: Bool -> Plugins.NameEnv [String] -> Plugins.CoreBind -> Bool
allowBind modAllowsRecursion anns = \case
  Plugins.NonRec {} -> True
  Plugins.Rec bs -> all (recursionAllowed modAllowsRecursion anns . fst) bs

recursionAllowed :: Bool -> Plugins.NameEnv [String] -> Plugins.Var -> Bool
recursionAllowed modAllowsRecursion anns var =
  let strAnns =
        Plugins.lookupWithDefaultUFM_Directly anns [] $ Plugins.getUnique var
   in (modAllowsRecursion || elem recursionAnnotation strAnns)
        && notElem noRecursionAnnotation strAnns

recursiveCallsForBind :: Plugins.Bind b -> [RecursionRecord b]
recursiveCallsForBind =
  let collectCalls v = addBindingReference v . collectRecursiveCalls
   in \case
        Plugins.NonRec v rhs -> collectCalls v rhs
        Plugins.Rec binds ->
          let nestedRecursion = foldMap (uncurry collectCalls) binds
           in maybe
                nestedRecursion
                (\bnds -> RecursionRecord [] (fst <$> bnds) : nestedRecursion)
                $ nonEmpty binds

-- | This collects all identifiable recursion points in an expression.
collectRecursiveCalls :: Plugins.Expr b -> [RecursionRecord b]
collectRecursiveCalls = \case
  Plugins.App f a -> collectRecursiveCalls f <> collectRecursiveCalls a
  Plugins.Case scrut _ _ alts ->
    collectRecursiveCalls scrut <> foldMap recursiveCallsForAlt alts
  Plugins.Cast e _ -> collectRecursiveCalls e
  Plugins.Coercion _ -> []
  Plugins.Lam _ body -> collectRecursiveCalls body
  Plugins.Let bind e -> recursiveCallsForBind bind <> collectRecursiveCalls e
  Plugins.Lit _ -> []
  Plugins.Tick _ body -> collectRecursiveCalls body
  Plugins.Type _ -> []
  Plugins.Var _ -> []

recursiveCallsForAlt :: Plugins.Alt b -> [RecursionRecord b]
recursiveCallsForAlt (Plugins.Alt _ _ rhs) = collectRecursiveCalls rhs