-- Copyright 2022 Google LLC
--
-- Use of this source code is governed by a BSD-style
-- license that can be found in the LICENSE file or at
-- https://developers.google.com/open-source/licenses/bsd

module Generalize (generalizeArgs, generalizeIxDict) where

import Control.Monad
import Data.Maybe (fromJust)

import Core
import Err
import PPrint
import Types.Core
import Inference
import IRVariants
import QueryType
import Name
import Subst
import Types.Primitives
import Types.Top

type RolePiBinder = WithAttrB RoleExpl CBinder
type RolePiBinders = Nest RolePiBinder

generalizeIxDict :: EnvReader m => CDict n -> m n (Generalized CoreIR CDict n)
generalizeIxDict dict = liftGeneralizerM do
  dict' <- sinkM dict
  dictTy <- return $ getType dict'
  dictTyGeneralized <- generalizeType dictTy
  liftEnvReaderM $ generalizeDict dictTyGeneralized dict'
{-# INLINE generalizeIxDict #-}

generalizeArgs ::EnvReader m => CorePiType n -> [Atom CoreIR n] -> m n (Generalized CoreIR (ListE CAtom) n)
generalizeArgs fTy argsTop = liftGeneralizerM $ runSubstReaderT idSubst do
  PairE (CorePiType _ expls bs _) (ListE argsTop') <- sinkM $ PairE fTy (ListE argsTop)
  ListE <$> go (zipAttrs expls bs) argsTop'
  where
    go :: Nest (WithAttrB Explicitness CBinder) i i' -> [Atom CoreIR n]
       -> SubstReaderT AtomSubstVal GeneralizerM i n [Atom CoreIR n]
    go (Nest (WithAttrB expl b) bs) (arg:args) = do
      ty' <- substM $ binderType b
      arg' <- liftSubstReaderT case (ty', expl) of
        (TyKind, _) -> toAtom <$> generalizeType (fromJust $ toMaybeType arg)
        (TyCon (DictTy _), Inferred Nothing (Synth _)) ->
          toAtom <$> generalizeDict ty' (fromJust $ toMaybeDict arg)
        _ -> isData ty' >>= \case
          True -> toAtom <$> emitGeneralizationParameter ty' arg
          False -> do
            -- Unlike in `inferRoles` in `Inference`, it's ok to have non-data,
            -- non-type, non-dict arguments (e.g. a function). We just don't
            -- generalize in that case.
            return arg
      args'' <- extendSubst (b@>SubstVal arg') $ go bs args
      return $ arg' : args''
    go Empty [] = return []
    go _ _ = error "zip error"
{-# INLINE generalizeArgs #-}

-- === generalizer monad plumbing ===

data GeneralizationEmission n l = GeneralizationEmission (Binder CoreIR n l) (Atom CoreIR n)
type GeneralizationEmissions = RNest GeneralizationEmission

newtype GeneralizerM n a = GeneralizerM {
  runGeneralizerM' :: DoubleInplaceT Env GeneralizationEmissions UnitB HardFailM n a }
  deriving (Functor, Applicative, Monad, MonadFail, Fallible, ScopeReader, EnvReader, EnvExtender)

liftGeneralizerM
  :: EnvReader m
  => (forall l. DExt n l => GeneralizerM l (e l))
  -> m n (Generalized CoreIR e n)
liftGeneralizerM cont = do
  env <- unsafeGetEnv
  Distinct <- getDistinct
  Abs emissions (DoubleInplaceTResult UnitB e) <- return $ runHardFail $
    runDoubleInplaceT env $ runGeneralizerM' cont
  let (bs, vals) = hoistGeneralizationVals (unRNest emissions)
  return (Abs bs e, vals)
  where
    -- OPTIMIZE: something not O(N^2)
    hoistGeneralizationVals :: Nest GeneralizationEmission n l -> (Nest (Binder CoreIR) n l, [Atom CoreIR n])
    hoistGeneralizationVals Empty = (Empty, [])
    hoistGeneralizationVals (Nest (GeneralizationEmission b val) bs) = do
      let (bs', vals) = hoistGeneralizationVals bs
      case hoist b (ListE vals) of
        HoistSuccess (ListE vals') -> (Nest b bs', val:vals')
        HoistFailure _ -> error "should't happen" -- when we do the generalization,
        -- the "local" values we emit never mention the new generalization binders.
        -- TODO: consider trying to encode this constraint using scope parameters.
{-# INLINE liftGeneralizerM #-}

-- XXX: the supplied type may be more general than the type of the atom!
emitGeneralizationParameter :: Type CoreIR n -> Atom CoreIR n -> GeneralizerM n (AtomVar CoreIR n)
emitGeneralizationParameter ty val = GeneralizerM do
  Abs b v <- return $ newName noHint
  let emission = Abs (RNest REmpty (GeneralizationEmission (b:>ty) val)) v
  emitDoubleInplaceTHoisted emission >>= \case
    -- This will hoist above binders appearing in types (e.g. pi binders, and
    -- dependent pair binders). As long as those variables are only used in
    -- DataParam roles, this hoisting should succeed.
    Nothing -> error $ "expected atom to be hoistable " ++ pprint val
    Just v' -> return $ AtomVar v' ty

-- === actual generalization traversal ===

-- Given a type (an Atom of type `Type`), abstracts over all data components
generalizeType :: Type CoreIR n -> GeneralizerM n (Type CoreIR n)
generalizeType ty = traverseTyParams ty \paramRole paramReqTy param -> case paramRole of
  TypeParam -> toAtom <$> generalizeType (fromJust $ toMaybeType param)
  DictParam -> toAtom <$> generalizeDict paramReqTy (fromJust $ toMaybeDict param)
  DataParam -> toAtom <$> emitGeneralizationParameter paramReqTy param

-- === role-aware type traversal ===

-- This traverses the type parameters, with knowledge of their roles and
-- expected types. It's used here for generalization, but it may also be useful
-- for other operations on types, like newtype stripping.

traverseTyParams
  :: forall m n. (EnvReader m, EnvExtender m)
  => CType n
  -> (forall l . DExt n l => ParamRole -> CType l -> CAtom l -> m l (CAtom l))
  -> m n (CType n)
traverseTyParams (StuckTy _ _) _ = error "shouldn't have StuckTy left"
traverseTyParams (TyCon ty) f = liftM TyCon $ getDistinct >>= \Distinct -> case ty of
  DictTy dictTy -> DictTy <$> case dictTy of
    DictType sn name params -> do
      Abs paramRoles UnitE <- getClassRoleBinders name
      params' <- traverseRoleBinders f paramRoles params
      return $ DictType sn name params'
    IxDictType   t -> IxDictType   <$> f' TypeParam TyKind t
    DataDictType t -> DataDictType <$> f' TypeParam TyKind t
  TabPi (TabPiType d (b:>iTy) resultTy) -> do
    iTy' <- f' TypeParam TyKind iTy
    let dictTy = toType $ IxDictType iTy'
    d' <- fromJust . toMaybeDict <$> f DictParam dictTy (toAtom d)
    withFreshBinder (getNameHint b) iTy' \(b':>_) -> do
      resultTy' <- applyRename (b@>binderName b') resultTy >>= (f' TypeParam TyKind)
      return $ TabPi $ TabPiType d' (b':>iTy') resultTy'
  BaseType b -> return $ BaseType b
  ProdType tys -> ProdType <$> forM tys \t -> f' TypeParam TyKind t
  RefType _ _ -> error "not implemented" -- how should we handle the ParamRole for the heap parameter?
  SumType  tys -> SumType  <$> forM tys \t -> f' TypeParam TyKind t
  TypeKind     -> return TypeKind
  HeapType     -> return HeapType
  NewtypeTyCon con -> NewtypeTyCon <$> case con of
    Nat -> return Nat
    Fin n -> Fin <$> f DataParam NatTy n
    EffectRowKind    -> return EffectRowKind
    UserADTType sn def (TyConParams infs params) -> do
      Abs roleBinders UnitE <- getDataDefRoleBinders def
      params' <- traverseRoleBinders f roleBinders params
      return $ UserADTType sn def $ TyConParams infs params'
  _ -> error $ "Not implemented: " ++ pprint ty
  where
    f' :: forall l . DExt n l => ParamRole -> CType l -> CType l -> m l (CType l)
    f' r t x = fromJust <$> toMaybeType <$> f r t (toAtom x)
{-# INLINE traverseTyParams #-}

traverseRoleBinders
  :: forall m n n'. EnvReader m
  => (forall l . DExt n l => ParamRole -> Type CoreIR l -> Atom CoreIR l -> m l (Atom CoreIR l))
  ->  RolePiBinders n n' -> [Atom CoreIR n] -> m n [Atom CoreIR n]
traverseRoleBinders f allBinders allParams =
  runSubstReaderT idSubst $ go allBinders allParams
  where
    go :: forall i i'. RolePiBinders i i' -> [Atom CoreIR n]
       -> SubstReaderT AtomSubstVal m i n [Atom CoreIR n]
    go Empty [] = return []
    go (Nest (WithAttrB (role, _) b) bs) (param:params) = do
      ty' <- substM $ binderType b
      Distinct <- getDistinct
      param' <- liftSubstReaderT $ f role ty' param
      params'' <- extendSubst (b@>SubstVal param') $ go bs params
      return $ param' : params''
    go _ _ = error "zip error"
{-# INLINE traverseRoleBinders #-}

getDataDefRoleBinders :: EnvReader m => TyConName n -> m n (Abs RolePiBinders UnitE n)
getDataDefRoleBinders def = do
  TyConDef _ attrs bs _ <- lookupTyCon def
  return $ Abs (zipAttrs attrs bs) UnitE
{-# INLINE getDataDefRoleBinders #-}

getClassRoleBinders :: EnvReader m => ClassName n -> m n (Abs RolePiBinders UnitE n)
getClassRoleBinders def = do
  ClassDef _ _ _ _ roleExpls bs _ _ <- lookupClassDef def
  return $ Abs (zipAttrs roleExpls bs) UnitE
{-# INLINE getClassRoleBinders #-}

-- === instances ===

instance GenericB GeneralizationEmission where
  type RepB GeneralizationEmission = BinderP (AtomNameC CoreIR) (PairE CType CAtom)
  fromB (GeneralizationEmission (b:>ty) x) = b :> PairE ty x
  {-# INLINE fromB #-}
  toB   (b :> PairE ty x) = GeneralizationEmission (b:>ty) x
  {-# INLINE toB #-}

instance RenameB GeneralizationEmission
instance HoistableB  GeneralizationEmission
instance ProvesExt   GeneralizationEmission
instance BindsNames  GeneralizationEmission
instance SinkableB   GeneralizationEmission

instance BindsEnv GeneralizationEmission where
  toEnvFrag (GeneralizationEmission b _) = toEnvFrag b
  {-# INLINE toEnvFrag #-}

instance ExtOutMap Env GeneralizationEmissions where
  extendOutMap bindings emissions =
    bindings `extendOutMap` toEnvFrag emissions
  {-# INLINE extendOutMap #-}