{-# LANGUAGE CPP #-}
{-# LANGUAGE GADTs #-}
{-# OPTIONS_GHC -fno-warn-incomplete-record-updates #-}

-- | A GHC plugin for a more Elm-like Haskell experience. It automatically
-- adds an unqualified import of the NriPrelude module, and qualified imports of
-- other base modules such as List and Maybe.
--
-- To use it make sure your project has @nri-prelude@ listed as a dependency,
-- then add the follwing ghc option to your cabal or package yaml file:
--
-- > -fplugin=NriPrelude.Plugin
module NriPrelude.Plugin
  ( plugin,
  )
where

-- Useful documentation
-- - Elm's default imports: https://package.elm-lang.org/packages/elm/core/latest/
-- - GHC user guide on compiler plugins: https://ghc.gitlab.haskell.org/ghc/doc/users_guide/extending_ghc.html#compiler-plugins
-- - Module providing API for creating plugins: https://hackage.haskell.org/package/ghc-lib-9.6.5.20240423/docs/GHC-Plugins.html

import Data.Function ((&))
import qualified Data.List
import qualified GHC.Hs
import qualified GHC.Parser.Annotation
import qualified GHC.Plugins
import NriPrelude.Plugin.GhcVersionDependent (setIDeclImplicit, withParsedResult)
import qualified Set
import Prelude

-- | adds an unqualified import of the NriPrelude module, and qualified imports of
-- other base modules such as List and Maybe.
--
-- To use it make sure your project has @nri-prelude@ listed as a dependency,
-- then add the follwing ghc option to your cabal or package yaml file:
--
-- > -fplugin=NriPrelude.Plugin
plugin :: GHC.Plugins.Plugin
plugin :: Plugin
plugin =
  Plugin
GHC.Plugins.defaultPlugin
    { GHC.Plugins.parsedResultAction = addImplicitImports,
      -- Let GHC know this plugin doesn't perform arbitrary IO. Given the same
      -- input file it will make the same changes. Without this GHC will
      -- recompile modules using this plugin every time which is expensive.
      GHC.Plugins.pluginRecompile = GHC.Plugins.purePlugin
    }

addImplicitImports ::
  [GHC.Plugins.CommandLineOption] ->
  GHC.Plugins.ModSummary ->
  GHC.Plugins.ParsedResult ->
  GHC.Plugins.Hsc GHC.Plugins.ParsedResult
addImplicitImports :: [CommandLineOption]
-> ModSummary -> ParsedResult -> Hsc ParsedResult
addImplicitImports [CommandLineOption]
_ ModSummary
_ ParsedResult
parsed =
  ParsedResult -> Hsc ParsedResult
forall a. a -> Hsc a
forall (f :: * -> *) a. Applicative f => a -> f a
Prelude.pure (ParsedResult -> Hsc ParsedResult)
-> ParsedResult -> Hsc ParsedResult
forall a b. (a -> b) -> a -> b
$
    ParsedResult -> (HsParsedModule -> HsParsedModule) -> ParsedResult
withParsedResult ParsedResult
parsed ((HsParsedModule -> HsParsedModule) -> ParsedResult)
-> (HsParsedModule -> HsParsedModule) -> ParsedResult
forall a b. (a -> b) -> a -> b
$ \HsParsedModule
parsed' ->
      HsParsedModule
parsed'
        { GHC.Hs.hpm_module =
            fmap addImportsWhenNotPath (GHC.Hs.hpm_module parsed')
        }
  where
    addImportsWhenNotPath :: HsModule GhcPs -> HsModule GhcPs
addImportsWhenNotPath HsModule GhcPs
hsModule =
      case (GenLocated SrcSpanAnnA ModuleName -> CommandLineOption)
-> Maybe (GenLocated SrcSpanAnnA ModuleName)
-> Maybe CommandLineOption
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap GenLocated SrcSpanAnnA ModuleName -> CommandLineOption
forall {l}. GenLocated l ModuleName -> CommandLineOption
unLocate (HsModule GhcPs -> Maybe (XRec GhcPs ModuleName)
forall p. HsModule p -> Maybe (XRec p ModuleName)
GHC.Hs.hsmodName HsModule GhcPs
hsModule) of
        Maybe CommandLineOption
Nothing -> HsModule GhcPs -> HsModule GhcPs
addImports HsModule GhcPs
hsModule
        Just CommandLineOption
modName ->
          if CommandLineOption -> CommandLineOption -> Bool
forall a. Eq a => [a] -> [a] -> Bool
Data.List.isPrefixOf CommandLineOption
"Paths_" CommandLineOption
modName
            then HsModule GhcPs
hsModule
            else HsModule GhcPs -> HsModule GhcPs
addImports HsModule GhcPs
hsModule

    addImports :: HsModule GhcPs -> HsModule GhcPs
addImports HsModule GhcPs
hsModule =
      HsModule GhcPs
hsModule
        { GHC.Hs.hsmodImports =
            -- Add default Elm-like imports when the user hasn't imported them
            -- explicitly yet, in order to avoid duplicate import warnings.
            GHC.Hs.hsmodImports hsModule
              ++ ( Set.diff extraImports (existingImports hsModule)
                     & Set.toList
                     & fmap
                       ( \Import
imp ->
                           case Import
imp of
                             Unqualified CommandLineOption
name -> CommandLineOption -> GenLocated SrcSpanAnnA (ImportDecl GhcPs)
forall {e}.
HasAnnotation e =>
CommandLineOption -> GenLocated e (ImportDecl GhcPs)
unqualified CommandLineOption
name
                             Qualified CommandLineOption
name -> CommandLineOption -> GenLocated SrcSpanAnnA (ImportDecl GhcPs)
forall {e}.
HasAnnotation e =>
CommandLineOption -> GenLocated e (ImportDecl GhcPs)
qualified CommandLineOption
name
                       )
                 )
        }

    existingImports :: HsModule p -> Set Import
existingImports HsModule p
hsModule =
      HsModule p -> [XRec p (ImportDecl p)]
forall p. HsModule p -> [LImportDecl p]
GHC.Hs.hsmodImports HsModule p
hsModule
        [GenLocated l (ImportDecl p)]
-> ([GenLocated l (ImportDecl p)] -> List Import) -> List Import
forall a b. a -> (a -> b) -> b
& (GenLocated l (ImportDecl p) -> Import)
-> [GenLocated l (ImportDecl p)] -> List Import
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap
          ( \(GHC.Plugins.L l
_ ImportDecl p
imp) ->
              case (ImportDecl p -> Bool
forall pass. ImportDecl pass -> Bool
isQualified ImportDecl p
imp, GenLocated l ModuleName -> CommandLineOption
forall {l}. GenLocated l ModuleName -> CommandLineOption
unLocate (ImportDecl p -> XRec p ModuleName
forall pass. ImportDecl pass -> XRec pass ModuleName
GHC.Hs.ideclName ImportDecl p
imp)) of
                (Bool
True, CommandLineOption
name) -> CommandLineOption -> Import
Qualified CommandLineOption
name
                (Bool
False, CommandLineOption
name) -> CommandLineOption -> Import
Unqualified CommandLineOption
name
          )
        List Import -> (List Import -> Set Import) -> Set Import
forall a b. a -> (a -> b) -> b
& List Import -> Set Import
forall comparable.
Ord comparable =>
List comparable -> Set comparable
Set.fromList

    unLocate :: GenLocated l ModuleName -> CommandLineOption
unLocate (GHC.Plugins.L l
_ ModuleName
x) = ModuleName -> CommandLineOption
GHC.Plugins.moduleNameString ModuleName
x

    unqualified :: CommandLineOption -> GenLocated e (ImportDecl GhcPs)
unqualified CommandLineOption
name =
      ImportDecl GhcPs -> GenLocated e (ImportDecl GhcPs)
forall e a. HasAnnotation e => a -> GenLocated e a
GHC.Parser.Annotation.noLocA (ModuleName -> ImportDecl GhcPs
GHC.Hs.simpleImportDecl (CommandLineOption -> ModuleName
GHC.Plugins.mkModuleName CommandLineOption
name))
        GenLocated e (ImportDecl GhcPs)
-> (GenLocated e (ImportDecl GhcPs)
    -> GenLocated e (ImportDecl GhcPs))
-> GenLocated e (ImportDecl GhcPs)
forall a b. a -> (a -> b) -> b
& (ImportDecl GhcPs -> ImportDecl GhcPs)
-> GenLocated e (ImportDecl GhcPs)
-> GenLocated e (ImportDecl GhcPs)
forall a b. (a -> b) -> GenLocated e a -> GenLocated e b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Bool -> ImportDecl GhcPs -> ImportDecl GhcPs
setIDeclImplicit Bool
True)
    qualified :: CommandLineOption -> GenLocated e (ImportDecl GhcPs)
qualified CommandLineOption
name =
      (ImportDecl GhcPs -> ImportDecl GhcPs)
-> GenLocated e (ImportDecl GhcPs)
-> GenLocated e (ImportDecl GhcPs)
forall a b. (a -> b) -> GenLocated e a -> GenLocated e b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\ImportDecl GhcPs
qual -> ImportDecl GhcPs
qual {GHC.Hs.ideclQualified = GHC.Hs.QualifiedPre}) (CommandLineOption -> GenLocated e (ImportDecl GhcPs)
forall {e}.
HasAnnotation e =>
CommandLineOption -> GenLocated e (ImportDecl GhcPs)
unqualified CommandLineOption
name)

-- There's more than one way to do a qualified import. See:
-- https://ghc.gitlab.haskell.org/ghc/doc/users_guide/exts/import_qualified_post.html

isQualified :: GHC.Hs.ImportDecl pass -> Bool
isQualified :: forall pass. ImportDecl pass -> Bool
isQualified ImportDecl pass
imp =
  case ImportDecl pass -> ImportDeclQualifiedStyle
forall pass. ImportDecl pass -> ImportDeclQualifiedStyle
GHC.Hs.ideclQualified ImportDecl pass
imp of
    ImportDeclQualifiedStyle
GHC.Hs.QualifiedPre -> Bool
True
    ImportDeclQualifiedStyle
GHC.Hs.QualifiedPost -> Bool
True
    ImportDeclQualifiedStyle
GHC.Hs.NotQualified -> Bool
False

data Import
  = Unqualified String
  | Qualified String
  deriving (Import -> Import -> Bool
(Import -> Import -> Bool)
-> (Import -> Import -> Bool) -> Eq Import
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Import -> Import -> Bool
== :: Import -> Import -> Bool
$c/= :: Import -> Import -> Bool
/= :: Import -> Import -> Bool
Eq, Eq Import
Eq Import =>
(Import -> Import -> Ordering)
-> (Import -> Import -> Bool)
-> (Import -> Import -> Bool)
-> (Import -> Import -> Bool)
-> (Import -> Import -> Bool)
-> (Import -> Import -> Import)
-> (Import -> Import -> Import)
-> Ord Import
Import -> Import -> Bool
Import -> Import -> Ordering
Import -> Import -> Import
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
$ccompare :: Import -> Import -> Ordering
compare :: Import -> Import -> Ordering
$c< :: Import -> Import -> Bool
< :: Import -> Import -> Bool
$c<= :: Import -> Import -> Bool
<= :: Import -> Import -> Bool
$c> :: Import -> Import -> Bool
> :: Import -> Import -> Bool
$c>= :: Import -> Import -> Bool
>= :: Import -> Import -> Bool
$cmax :: Import -> Import -> Import
max :: Import -> Import -> Import
$cmin :: Import -> Import -> Import
min :: Import -> Import -> Import
Ord)

-- taken from https://package.elm-lang.org/packages/elm/core/latest/
extraImports :: Set.Set Import
extraImports :: Set Import
extraImports =
  List Import -> Set Import
forall comparable.
Ord comparable =>
List comparable -> Set comparable
Set.fromList
    [ CommandLineOption -> Import
Unqualified CommandLineOption
"NriPrelude", -- Elm exports types from withi these modules. We re-export them from NriPrelude. Same effect.
      CommandLineOption -> Import
Qualified CommandLineOption
"Basics",
      CommandLineOption -> Import
Qualified CommandLineOption
"Char",
      CommandLineOption -> Import
Qualified CommandLineOption
"Debug",
      CommandLineOption -> Import
Qualified CommandLineOption
"List",
      CommandLineOption -> Import
Qualified CommandLineOption
"Maybe",
      CommandLineOption -> Import
Qualified CommandLineOption
"Platform",
      CommandLineOption -> Import
Qualified CommandLineOption
"Result",
      CommandLineOption -> Import
Qualified CommandLineOption
"Text", -- equivalent to Elm's String
      CommandLineOption -> Import
Qualified CommandLineOption
"Tuple",
      -- Additionally Task and Log because we use them everywhere
      CommandLineOption -> Import
Qualified CommandLineOption
"Log",
      CommandLineOption -> Import
Qualified CommandLineOption
"Task"
    ]