{-# LANGUAGE PatternGuards #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE TypeFamilies #-}

-- | See <https://github.com/ezyang/ghc-proposals/blob/backpack/proposals/0000-backpack.rst>
module Distribution.Backpack.ReadyComponent
  ( ReadyComponent (..)
  , InstantiatedComponent (..)
  , IndefiniteComponent (..)
  , rc_depends
  , rc_uid
  , rc_pkgid
  , dispReadyComponent
  , toReadyComponents
  ) where

import Distribution.Compat.Prelude hiding ((<>))
import Prelude ()

import Distribution.Backpack
import Distribution.Backpack.LinkedComponent
import Distribution.Backpack.ModuleShape

import Distribution.Compat.Graph (IsNode (..))
import Distribution.Types.AnnotatedId
import Distribution.Types.Component
import Distribution.Types.ComponentId
import Distribution.Types.ComponentInclude
import Distribution.Types.ComponentName
import Distribution.Types.Library
import Distribution.Types.LibraryName
import Distribution.Types.Module
import Distribution.Types.ModuleRenaming
import Distribution.Types.MungedPackageId
import Distribution.Types.MungedPackageName
import Distribution.Types.PackageId
import Distribution.Types.PackageName.Magic
import Distribution.Types.UnitId

import Distribution.ModuleName
import Distribution.Package
import Distribution.Simple.Utils

import Control.Monad
import qualified Data.Map as Map
import qualified Data.Set as Set
import Text.PrettyPrint

import Distribution.Pretty
import Distribution.Version

-- | A 'ReadyComponent' is one that we can actually generate build
-- products for.  We have a ready component for the typecheck-only
-- products of every indefinite package, as well as a ready component
-- for every way these packages can be fully instantiated.
data ReadyComponent = ReadyComponent
  { rc_ann_id :: AnnotatedId UnitId
  , rc_open_uid :: OpenUnitId
  -- ^ The 'OpenUnitId' for this package.  At the moment, this
  -- is used in only one case, which is to determine if an
  -- export is of a module from this library (indefinite
  -- libraries record these exports as 'OpenModule');
  -- 'rc_open_uid' can be conveniently used to test for
  -- equality, whereas 'UnitId' cannot always be used in this
  -- case.
  , rc_cid :: ComponentId
  -- ^ Corresponds to 'lc_cid'.  Invariant: if 'rc_open_uid'
  -- records a 'ComponentId', it coincides with this one.
  , rc_component :: Component
  -- ^ Corresponds to 'lc_component'.
  , rc_exe_deps :: [AnnotatedId UnitId]
  -- ^ Corresponds to 'lc_exe_deps'.
  -- Build-tools don't participate in mix-in linking.
  -- (but what if they could?)
  , rc_public :: Bool
  -- ^ Corresponds to 'lc_public'.
  , rc_i :: Either IndefiniteComponent InstantiatedComponent
  -- ^ Extra metadata depending on whether or not this is an
  -- indefinite library (typechecked only) or an instantiated
  -- component (can be compiled).
  }

-- | The final, string 'UnitId' that will uniquely identify
-- the compilation products of this component.
rc_uid :: ReadyComponent -> UnitId
rc_uid = ann_id . rc_ann_id

-- | Corresponds to 'lc_pkgid'.
rc_pkgid :: ReadyComponent -> PackageId
rc_pkgid = ann_pid . rc_ann_id

-- | An 'InstantiatedComponent' is a library which is fully instantiated
-- (or, possibly, has no requirements at all.)
data InstantiatedComponent = InstantiatedComponent
  { instc_insts :: [(ModuleName, Module)]
  -- ^ How this library was instantiated.
  , instc_insts_deps :: [(UnitId, MungedPackageId)]
  -- ^ Dependencies induced by 'instc_insts'.  These are recorded
  -- here because there isn't a convenient way otherwise to get
  -- the 'PackageId' we need to fill 'componentPackageDeps' as needed.
  , instc_provides :: Map ModuleName Module
  -- ^ The modules exported/reexported by this library.
  , instc_includes :: [ComponentInclude DefUnitId ModuleRenaming]
  -- ^ The dependencies which need to be passed to the compiler
  -- to bring modules into scope.  These always refer to installed
  -- fully instantiated libraries.
  }

-- | An 'IndefiniteComponent' is a library with requirements
-- which we will typecheck only.
data IndefiniteComponent = IndefiniteComponent
  { indefc_requires :: [ModuleName]
  -- ^ The requirements of the library.
  , indefc_provides :: Map ModuleName OpenModule
  -- ^ The modules exported/reexported by this library.
  , indefc_includes :: [ComponentInclude OpenUnitId ModuleRenaming]
  -- ^ The dependencies which need to be passed to the compiler
  -- to bring modules into scope.  These are 'OpenUnitId' because
  -- these may refer to partially instantiated libraries.
  }

-- | Compute the dependencies of a 'ReadyComponent' that should
-- be recorded in the @depends@ field of 'InstalledPackageInfo'.
rc_depends :: ReadyComponent -> [(UnitId, MungedPackageId)]
rc_depends rc = ordNub $
  case rc_i rc of
    Left indefc ->
      map
        (\ci -> (abstractUnitId $ ci_id ci, toMungedPackageId ci))
        (indefc_includes indefc)
    Right instc ->
      map
        (\ci -> (unDefUnitId $ ci_id ci, toMungedPackageId ci))
        (instc_includes instc)
        ++ instc_insts_deps instc
  where
    toMungedPackageId :: Pretty id => ComponentInclude id rn -> MungedPackageId
    toMungedPackageId ci =
      computeCompatPackageId
        (ci_pkgid ci)
        ( case ci_cname ci of
            CLibName name -> name
            _ ->
              error $
                prettyShow (rc_cid rc)
                  ++ " depends on non-library "
                  ++ prettyShow (ci_id ci)
        )

-- | Get the 'MungedPackageId' of a 'ReadyComponent' IF it is
-- a library.
rc_munged_id :: ReadyComponent -> MungedPackageId
rc_munged_id rc =
  computeCompatPackageId
    (rc_pkgid rc)
    ( case rc_component rc of
        CLib lib -> libName lib
        _ -> error "rc_munged_id: not library"
    )

instance Package ReadyComponent where
  packageId = rc_pkgid

instance HasUnitId ReadyComponent where
  installedUnitId = rc_uid

instance IsNode ReadyComponent where
  type Key ReadyComponent = UnitId
  nodeKey = rc_uid
  nodeNeighbors rc =
    ( case rc_i rc of
        Right inst
          | [] <- instc_insts inst ->
              []
          | otherwise ->
              [newSimpleUnitId (rc_cid rc)]
        _ -> []
    )
      ++ ordNub (map fst (rc_depends rc))
      ++ map ann_id (rc_exe_deps rc)

dispReadyComponent :: ReadyComponent -> Doc
dispReadyComponent rc =
  hang
    ( text
        ( case rc_i rc of
            Left _ -> "indefinite"
            Right _ -> "definite"
        )
        <+> pretty (nodeKey rc)
        {- <+> dispModSubst (Map.fromList (lc_insts lc)) -}
    )
    4
    $ vcat
      [ text "depends" <+> pretty uid
      | uid <- nodeNeighbors rc
      ]

-- | The state of 'InstM'; a mapping from 'UnitId's to their
-- ready component, or @Nothing@ if its an external
-- component which we don't know how to build.
type InstS = Map UnitId (Maybe ReadyComponent)

-- | A state monad for doing instantiations (can't use actual
-- State because that would be an extra dependency.)
newtype InstM a = InstM {runInstM :: InstS -> (a, InstS)}

instance Functor InstM where
  fmap f (InstM m) = InstM $ \s ->
    let (x, s') = m s
     in (f x, s')

instance Applicative InstM where
  pure a = InstM (a,)
  InstM f <*> InstM x = InstM $ \s ->
    let (f', s') = f s
        (x', s'') = x s'
     in (f' x', s'')

instance Monad InstM where
  return = pure
  InstM m >>= f = InstM $ \s ->
    let (x, s') = m s
     in runInstM (f x) s'

-- | Given a list of 'LinkedComponent's, expand the module graph
-- so that we have an instantiated graph containing all of the
-- instantiated components we need to build.
--
-- Instantiation intuitively follows the following algorithm:
--
--      instantiate a definite unit id p[S]:
--          recursively instantiate each module M in S
--          recursively instantiate modules exported by this unit
--          recursively instantiate dependencies substituted by S
--
-- The implementation is a bit more involved to memoize instantiation
-- if we have done it already.
--
-- We also call 'improveUnitId' during this process, so that fully
-- instantiated components are given 'HashedUnitId'.
toReadyComponents
  :: Map UnitId MungedPackageId
  -> Map ModuleName Module -- subst for the public component
  -> [LinkedComponent]
  -> [ReadyComponent]
toReadyComponents pid_map subst0 comps =
  catMaybes (Map.elems ready_map)
  where
    cmap = Map.fromList [(lc_cid lc, lc) | lc <- comps]

    instantiateUnitId
      :: ComponentId
      -> Map ModuleName Module
      -> InstM DefUnitId
    instantiateUnitId cid insts = InstM $ \s ->
      case Map.lookup uid s of
        Nothing ->
          -- Knot tied
          let (r, s') =
                runInstM
                  (instantiateComponent uid cid insts)
                  (Map.insert uid r s)
           in (def_uid, Map.insert uid r s')
        Just _ -> (def_uid, s)
      where
        -- The mkDefUnitId here indicates that we assume
        -- that Cabal handles unit id hash allocation.
        -- Good thing about hashing here: map is only on string.
        -- Bad thing: have to repeatedly hash.
        def_uid = mkDefUnitId cid insts
        uid = unDefUnitId def_uid

    instantiateComponent
      :: UnitId
      -> ComponentId
      -> Map ModuleName Module
      -> InstM (Maybe ReadyComponent)
    instantiateComponent uid cid insts
      | Just lc <- Map.lookup cid cmap = do
          provides <- traverse (substModule insts) (modShapeProvides (lc_shape lc))
          -- NB: lc_sig_includes is omitted here, because we don't
          -- need them to build
          includes <- forM (lc_includes lc) $ \ci -> do
            uid' <- substUnitId insts (ci_id ci)
            return ci{ci_ann_id = fmap (const uid') (ci_ann_id ci)}
          exe_deps <- traverse (substExeDep insts) (lc_exe_deps lc)
          s <- InstM $ \s -> (s, s)
          let getDep (Module dep_def_uid _)
                | let dep_uid = unDefUnitId dep_def_uid =
                    -- Lose DefUnitId invariant for rc_depends
                    [
                      ( dep_uid
                      , fromMaybe err_pid $
                          Map.lookup dep_uid pid_map
                            <|> fmap rc_munged_id (join (Map.lookup dep_uid s))
                      )
                    ]
                where
                  err_pid =
                    MungedPackageId
                      (MungedPackageName nonExistentPackageThisIsCabalBug LMainLibName)
                      (mkVersion [0])
              instc =
                InstantiatedComponent
                  { instc_insts = Map.toList insts
                  , instc_insts_deps = concatMap getDep (Map.elems insts)
                  , instc_provides = provides
                  , instc_includes = includes
                  -- NB: there is no dependency on the
                  -- indefinite version of this instantiated package here,
                  -- as (1) it doesn't go in depends in the
                  -- IPI: it's not a run time dep, and (2)
                  -- we don't have to tell GHC about it, it
                  -- will match up the ComponentId
                  -- automatically
                  }
          return $
            Just
              ReadyComponent
                { rc_ann_id = (lc_ann_id lc){ann_id = uid}
                , rc_open_uid = DefiniteUnitId (unsafeMkDefUnitId uid)
                , rc_cid = lc_cid lc
                , rc_component = lc_component lc
                , rc_exe_deps = exe_deps
                , rc_public = lc_public lc
                , rc_i = Right instc
                }
      | otherwise = return Nothing

    substUnitId :: Map ModuleName Module -> OpenUnitId -> InstM DefUnitId
    substUnitId _ (DefiniteUnitId uid) =
      return uid
    substUnitId subst (IndefFullUnitId cid insts) = do
      insts' <- substSubst subst insts
      instantiateUnitId cid insts'

    -- NB: NOT composition
    substSubst
      :: Map ModuleName Module
      -> Map ModuleName OpenModule
      -> InstM (Map ModuleName Module)
    substSubst subst insts = traverse (substModule subst) insts

    substModule :: Map ModuleName Module -> OpenModule -> InstM Module
    substModule subst (OpenModuleVar mod_name)
      | Just m <- Map.lookup mod_name subst = return m
      | otherwise = error "substModule: non-closing substitution"
    substModule subst (OpenModule uid mod_name) = do
      uid' <- substUnitId subst uid
      return (Module uid' mod_name)

    substExeDep
      :: Map ModuleName Module
      -> AnnotatedId OpenUnitId
      -> InstM (AnnotatedId UnitId)
    substExeDep insts exe_aid = do
      exe_uid' <- substUnitId insts (ann_id exe_aid)
      return exe_aid{ann_id = unDefUnitId exe_uid'}

    indefiniteUnitId :: ComponentId -> InstM UnitId
    indefiniteUnitId cid = do
      let uid = newSimpleUnitId cid
      r <- indefiniteComponent uid cid
      InstM $ \s -> (uid, Map.insert uid r s)

    indefiniteComponent :: UnitId -> ComponentId -> InstM (Maybe ReadyComponent)
    indefiniteComponent uid cid
      | Just lc <- Map.lookup cid cmap = do
          -- We're going to process includes, in case some of them
          -- are fully definite even without any substitution.  We
          -- want to build those too; see #5634.
          inst_includes <- forM (lc_includes lc) $ \ci ->
            if Set.null (openUnitIdFreeHoles (ci_id ci))
              then do
                uid' <- substUnitId Map.empty (ci_id ci)
                return $ ci{ci_ann_id = fmap (const (DefiniteUnitId uid')) (ci_ann_id ci)}
              else return ci
          exe_deps <- traverse (substExeDep Map.empty) (lc_exe_deps lc)
          let indefc =
                IndefiniteComponent
                  { indefc_requires = map fst (lc_insts lc)
                  , indefc_provides = modShapeProvides (lc_shape lc)
                  , indefc_includes = inst_includes ++ lc_sig_includes lc
                  }
          return $
            Just
              ReadyComponent
                { rc_ann_id = (lc_ann_id lc){ann_id = uid}
                , rc_cid = lc_cid lc
                , rc_open_uid = lc_uid lc
                , rc_component = lc_component lc
                , -- It's always fully built
                  rc_exe_deps = exe_deps
                , rc_public = lc_public lc
                , rc_i = Left indefc
                }
      | otherwise = return Nothing

    ready_map = snd $ runInstM work Map.empty

    work
      -- Top-level instantiation per subst0
      | not (Map.null subst0)
      , [lc] <- filter lc_public (Map.elems cmap) =
          do
            _ <- instantiateUnitId (lc_cid lc) subst0
            return ()
      | otherwise =
          forM_ (Map.elems cmap) $ \lc ->
            if null (lc_insts lc)
              then instantiateUnitId (lc_cid lc) Map.empty >> return ()
              else indefiniteUnitId (lc_cid lc) >> return ()
