content: {-# LANGUAGE GeneralizedNewtypeDeriving, OverloadedStrings, NamedFieldPuns #-}
  {-# LANGUAGE GADTs, RankNTypes, ScopedTypeVariables, DataKinds, TupleSections #-}
  {-# LANGUAGE FlexibleContexts, KindSignatures, PolyKinds #-}
  {-# LANGUAGE RecordWildCards #-} -- for dealing with TCDecl and existential k
  
  {- |
  Specialise rules s.t. polymorphic rules and those which have
  grammar arguments are removed.  This also has the effect of
  removing the information about recursive grouping (FIXME: we could
  preserve it).
  
  Consider a declaration:
  
      f as x P = e      -- `as` are the type parameters to the function
  
  for each call site, A, `f ts x (Q y)` we generate a new function:
  
      f_A x y = e [ts/as] [(Q y/P]
  
  we also try to reuse instances, so if there are some other call sites,
  for examle:
    B: f ts y (Q z)
    C: f ts y (Q y)
  
  we are going to just reuse `f_A` like this:
    B: f_A y z
    C: f_A y y
  
  Type arguments are just compared for equality, while other arguments
  are unified.  See 'Specialise.Unfiy' for details.
  -}
  
  
  module Daedalus.Specialise (specialise, regroup) where
  
  import Control.Monad
  import Data.Map (Map)
  import qualified Data.Map as Map
  import qualified Data.Set as Set
  import Data.Semigroup (First(..))
  import MonadLib
  
  import Data.Parameterized.Some
  
  import Daedalus.Panic
  import Daedalus.SourceRange
  import Daedalus.PP
  import Daedalus.Pass
  
  
  
  import Daedalus.AST(nameScopeAsModScope)
  import Daedalus.Type.AST
  import Daedalus.Type.Traverse
  
  import Daedalus.Type.Free
  import Daedalus.Specialise.Monad
  import Daedalus.Specialise.Unify
  import Data.Maybe (isNothing)
  
  -- -----------------------------------------------------------------------------
  -- Top level driver
  
  
  regroup :: [TCDecl SourceRange] -> Map ModuleName [Rec (TCDecl SourceRange)]
  regroup = fmap reverse . foldr addR Map.empty . topoOrder
    where
    owner d = case nameScopedIdent (tcDeclName d) of
                ModScope m _ -> m
                _ -> panic "regroup" [ "Declaration is not ModScope" ]
  
    addR r mp =
      case r of
        NonRec d -> Map.insertWith (++) (owner d) [r] mp
        MutRec ds -> case map owner ds of
                       x : xs | all (== x) xs ->
                                Map.insertWith (++) x [r] mp
                       _ -> panic "recroup.addR" ["Oops"]
  
  
  -- | This assumes that the declratations are in dependency order.
  specialise :: [Name] -> [Rec (TCDecl SourceRange)]
                -> PassM (Either String [TCDecl SourceRange])
  specialise ruleRoots decls =
    runPApplyM ruleRoots (concat . reverse <$> mapM go (reverse decls))
    where
      -- First we find if we need to generate partial applications.  If we
      -- do so, we can discard the input tdecl (FIXME: I think?), the
      -- reasoning being that specialised decls are problematic otherwise.
      go (NonRec d) = do
        insts <- getPendingSpecs [tcDeclName d]
        seen  <- seenRule (tcDeclName d)
        ds    <- case Map.lookup (tcDeclName d) insts of
                   Just is        -> mapM (flip apInst d) is -- forget d
                   Nothing | seen -> pure [d]
                   Nothing        -> pure []
        mapM specialiseOne ds
  
      -- We treat each inst. req. independently; the unify stuff will
      -- make sure it doesn't explode.  We check for unsupported
      -- recursion by checking for new recursive reqs. after completing
      -- all decls in the ds.
      go (MutRec ds) = do
        insts <- getPendingSpecs (map tcDeclName ds)
        -- We know that any recursive element we have seen doesn't have
        -- grammar args, so is a reasonably place to start if we have no
        -- specialisations.  Note that we may have specialisations that
        -- are internal to the recursive group.
        (seen, unseen) <- partitionM (seenRule . tcDeclName) ds
        let roots = [ Map.singleton n [i] | (n, is) <- Map.toList insts, i <- is ]
        case (roots, seen) of
          ([], [])      -> pure []
          ([], _ )      -> goOneRoot (seen ++ unseen) mempty
          _             -> concat <$> mapM (goOneRoot ds) roots
  
      goOneRoot :: [TCDecl SourceRange] ->
                   Map Name [Instantiation] -> PApplyM [TCDecl SourceRange]
      goOneRoot ds todo = do
        rs <- goOne ds todo
        insts <- getPendingSpecs (map tcDeclName ds)
        if not (Map.null insts)
          then raise ("Incompatible recursion detected for " ++ show (ppError insts))
          else return rs
  
      ppError insts =
        punctuate ", " $ map (\(k, is) -> pp k <+> vcat (map pp is)) (Map.toList insts)
  
      -- The idea here is that we process each decl exactly once, so
      -- that if we see the decl again (after we are done), then
      -- something is wrong with the input.
      goOne :: [TCDecl SourceRange] ->
               Map Name [Instantiation] -> PApplyM [TCDecl SourceRange]
      goOne [] _    = pure []
      goOne ds@(d : rest) todo =
        case Map.minViewWithKey todo of
          -- Fall back to just processing the first decl.
          Nothing ->
              do d'      <- specialiseOne d
                 newTodo <- getPendingSpecs (map tcDeclName rest)
                 (d' :) <$> goOne rest newTodo
  
          Just ((n, [inst]), todoRest)
            | (ds', d' : ds'') <- break (\di -> tcDeclName di == n) ds ->
              do let newds = ds' ++ ds''
                 -- let d' = apInst inst d
                 d'' <- specialiseOne =<< apInst inst d'
                 newTodo <- getPendingSpecs (map tcDeclName newds)
                 (d'' :) <$> goOne newds (Map.unionWith (++) todoRest newTodo)
  
          Just ((n, _ : _ : _), _) ->
            raise ("Multiple instantiations requested for " ++ show (pp n))
          Just ((_, []), _) -> panic "Empty instantiation list" []
          _ -> panic "Impossible" []
  
  
  -- This function traverses a term and replaces all problematic
  -- function calls by specialised versions
  specialiseOne :: TCDecl SourceRange -> PApplyM (TCDecl SourceRange)
  specialiseOne TCDecl {..}
    | not (null tcDeclTyParams) = panic "specialiseOne"
                                        ["Specializing a poly function"]
    | otherwise =
    case tcDeclDef of
      ExternDecl _ -> pure TCDecl { .. }
      Defined d -> do tdef <- go d
                      pure (TCDecl { tcDeclDef = Defined tdef, .. })
    where
      go :: forall k'. TC SourceRange k' -> PApplyM (TC SourceRange k')
      go (TC v) = TC <$> traverse go' v
  
      go' :: forall k'. TCF SourceRange k' -> PApplyM (TCF SourceRange k')
      go' texpr =
        case texpr of
          -- FIXME: maype specialise simple recursive case?
          TCCall n ts as -> do
            as' <- mapM (traverseArg go) as
            let m = fst (nameScopeAsModScope tcDeclName)
            specialiseCall m n ts as'
          x -> traverseTCF go x
  
  partitionM :: Monad m => (a -> m Bool) -> [a] -> m ([a], [a])
  partitionM p xs = do
    (yess, nos) <- foldM go ([], []) xs
    pure (reverse yess, reverse nos)
    where
      go (yes, no) x = do r <- p x
                          pure $ if r then (x : yes, no) else (yes, x : no)
  
  
  
  -- -----------------------------------------------------------------------------
  -- Specialisation policy
  
  -- This is the main policy function --- this determines how and when a
  -- function call necessitates a specialised version.  Returns a new
  -- call if required.
  specialiseCall ::
    ModuleName        {- ^ Name of the module containing the call -} ->
    TCName k          {- ^ Call this function -} ->
    [Type]            {- ^ With these types -} ->
    [Arg SourceRange] {- ^ And these concrete arguments -} ->
    PApplyM (TCF SourceRange k)
  
  -- No specialisation required if there are no type args, and no grammar args.
  specialiseCall m n ts args
    | [] <- ts, all isNothing probArgs = do
        addSeenRule (tcName n)
        pure (TCCall n [] args)
    | otherwise = requestSpec m n ts probArgs args
    where
      probArgs = map probArg args
  
      -- If it is a partially applied function, we inline.
      probArg arg | Type (TFun _ _) <- typeOf arg = Just arg
      -- Any non-function typed value is left alone
      probArg (ValArg _) = Nothing
      -- Anything else is inlined.
      probArg arg        = Just arg
  
  -- -----------------------------------------------------------------------------
  
  {- Request a specialisation
  
  We want to specialise a call, this checks to see if it unifies with
  an existing spec. request.
  -}
  requestSpec ::
    ModuleName            {- ^ Name of the module containing the call -} ->
    TCName k              {- ^ Specialize this -} ->
    [Type]                {- ^ Using these type arguments -} ->
    [Maybe (Arg SourceRange)]
                          {- ^ And these arguments: Nothing = leave as arg -}->
    [Arg SourceRange]     {- ^ All original arguments -} ->
    PApplyM (TCF SourceRange k)
  requestSpec m tnm ts args origArgs = do
    rs <- lookupRequestedSpecs (tcName tnm)
    case rs of
      Just insts | Just (First call) <- foldMap findUnifier insts
                   -> pure call
      _ -> do nm' <- addSpecRequest m (tcName tnm) ts newPs args
              pure (mkCall nm' newPs)
    where
      newPs = map getValue (Set.toList (tcFree args))
  
      getValue :: Some TCName -> TCName Value
      getValue (Some tnm'@(TCName { tcNameCtx = AValue })) = tnm'
      getValue (Some tnm') =
          panic "requestSpec"
            [ "Saw a non-Value free variable: " ++ show (pp tnm')
            , "Requesting " ++ show (pp tnm)
            , "Args: " ++ show (hsep $ map ppA args)
            ]
  
      ppA Nothing = text "_"
      ppA (Just v) = pp v
  
      findUnifier inst
        | ts == instTys inst
        , Right unifier <- unify args (instArgs inst)
        = let params = map (apUnifier unifier) (instNewParams inst)
          in Just (First (mkCall (instNewName inst) params))
  
      findUnifier _ = Nothing
  
      remainingArgs = [ a | (a, Nothing) <- zip origArgs args ]
  
      -- We don't pass any type args, as the target should be monomorphic
      mkCall nm' params =
        TCCall (tnm { tcName = nm' }) []
               (map (ValArg . syntheticTC . TCVar) params ++ remainingArgs)
  
  syntheticTC :: TCF SourceRange k -> TC SourceRange k
  syntheticTC = annotExpr synthetic
  
  
  
  