content: -- | This module implements the generic deriving elaboration that takes place during desugaring.
  module Language.PureScript.Sugar.TypeClasses.Deriving (deriveInstances) where
  
  import Prelude
  import Protolude (note)
  
  import Control.Monad.Error.Class (MonadError(..))
  import Control.Monad.Supply.Class (MonadSupply)
  import Data.List (foldl', find, unzip5)
  import Language.PureScript.AST (Binder(..), CaseAlternative(..), DataConstructorDeclaration(..), Declaration(..), Expr(..), pattern MkUnguarded, Module(..), SourceSpan(..), TypeInstanceBody(..), pattern ValueDecl)
  import Language.PureScript.AST.Utils (UnwrappedTypeConstructor(..), lamCase, unguarded, unwrapTypeConstructor)
  import Language.PureScript.Constants.Libs qualified as Libs
  import Language.PureScript.Crash (internalError)
  import Language.PureScript.Environment (DataDeclType(..), NameKind(..))
  import Language.PureScript.Errors (MultipleErrors, SimpleErrorMessage(..), errorMessage')
  import Language.PureScript.Names (pattern ByNullSourcePos, Ident(..), ModuleName, ProperName(..), ProperNameType(..), Qualified(..), QualifiedBy(..), freshIdent)
  import Language.PureScript.PSString (mkString)
  import Language.PureScript.Types (SourceType, Type(..), WildcardData(..), replaceAllTypeVars, srcTypeApp, srcTypeConstructor, srcTypeLevelString)
  import Language.PureScript.TypeChecker (checkNewtype)
  
  -- | Elaborates deriving instance declarations by code generation.
  deriveInstances
    :: forall m
     . (MonadError MultipleErrors m, MonadSupply m)
    => Module
    -> m Module
  deriveInstances (Module ss coms mn ds exts) =
      Module ss coms mn <$> mapM (deriveInstance mn ds) ds <*> pure exts
  
  -- | Takes a declaration, and if the declaration is a deriving TypeInstanceDeclaration,
  -- elaborates that into an instance declaration via code generation.
  --
  -- More instance deriving happens during type checking. The instances
  -- derived here are special for two reasons:
  -- * they depend only on the structure of the data, not types; and
  -- * they expect wildcard types from the user and generate type expressions
  --   to replace them.
  --
  deriveInstance
    :: forall m
     . (MonadError MultipleErrors m, MonadSupply m)
    => ModuleName
    -> [Declaration]
    -> Declaration
    -> m Declaration
  deriveInstance mn ds decl =
    case decl of
      TypeInstanceDeclaration sa@(ss, _) na ch idx nm deps className tys DerivedInstance -> let
        binaryWildcardClass :: (Declaration -> [SourceType] -> m ([Declaration], SourceType)) -> m Declaration
        binaryWildcardClass f = case tys of
          [ty1, ty2] -> case unwrapTypeConstructor ty1 of
            Just UnwrappedTypeConstructor{..} | mn == utcModuleName -> do
              checkIsWildcard ss utcTyCon ty2
              tyConDecl <- findTypeDecl ss utcTyCon ds
              (members, ty2') <- f tyConDecl utcArgs
              pure $ TypeInstanceDeclaration sa na ch idx nm deps className [ty1, ty2'] (ExplicitInstance members)
            _ -> throwError . errorMessage' ss $ ExpectedTypeConstructor className tys ty1
          _ -> throwError . errorMessage' ss $ InvalidDerivedInstance className tys 2
  
        in case className of
          Libs.Generic -> binaryWildcardClass (deriveGenericRep ss mn)
          Libs.Newtype -> binaryWildcardClass deriveNewtype
          _ -> pure decl
      _ -> pure decl
  
  deriveGenericRep
    :: forall m
     . (MonadError MultipleErrors m, MonadSupply m)
    => SourceSpan
    -> ModuleName
    -> Declaration
    -> [SourceType]
    -> m ([Declaration], SourceType)
  deriveGenericRep ss mn tyCon tyConArgs =
    case tyCon of
      DataDeclaration (ss', _) _ _ args dctors -> do
        x <- freshIdent "x"
        (reps, to, from) <- unzip3 <$> traverse makeInst dctors
        let rep = toRepTy reps
            inst | null reps =
                     -- If there are no cases, spin
                     [ ValueDecl (ss', []) (Ident "to") Public [] $ unguarded $
                        lamCase x
                          [ CaseAlternative
                              [NullBinder]
                              (unguarded (App (Var ss Libs.I_to) (Var ss' (Qualified ByNullSourcePos x))))
                          ]
                     , ValueDecl (ss', []) (Ident "from") Public [] $ unguarded $
                        lamCase x
                          [ CaseAlternative
                              [NullBinder]
                              (unguarded (App (Var ss Libs.I_from) (Var ss' (Qualified ByNullSourcePos x))))
                          ]
                     ]
                 | otherwise =
                     [ ValueDecl (ss', []) (Ident "to") Public [] $ unguarded $
                         lamCase x (zipWith ($) (map underBinder (sumBinders (length dctors))) to)
                     , ValueDecl (ss', []) (Ident "from") Public [] $ unguarded $
                         lamCase x (zipWith ($) (map underExpr (sumExprs (length dctors))) from)
                     ]
  
            subst = zipWith ((,) . fst) args tyConArgs
        return (inst, replaceAllTypeVars subst rep)
      _ -> internalError "deriveGenericRep: expected DataDeclaration"
  
      where
  
      select :: (a -> a) -> (a -> a) -> Int -> [a -> a]
      select _ _ 0 = []
      select _ _ 1 = [id]
      select l r n = take (n - 1) (iterate (r .) l) ++ [compN (n - 1) r]
  
      sumBinders :: Int -> [Binder -> Binder]
      sumBinders = select (ConstructorBinder ss Libs.C_Inl . pure)
                          (ConstructorBinder ss Libs.C_Inr . pure)
  
      sumExprs :: Int -> [Expr -> Expr]
      sumExprs = select (App (Constructor ss Libs.C_Inl))
                        (App (Constructor ss Libs.C_Inr))
  
      compN :: Int -> (a -> a) -> a -> a
      compN 0 _ = id
      compN n f = f . compN (n - 1) f
  
      makeInst
        :: DataConstructorDeclaration
        -> m (SourceType, CaseAlternative, CaseAlternative)
      makeInst (DataConstructorDeclaration _ ctorName args) = do
          let args' = map snd args
          (ctorTy, matchProduct, ctorArgs, matchCtor, mkProduct) <- makeProduct args'
          return ( srcTypeApp (srcTypeApp (srcTypeConstructor Libs.Constructor)
                                    (srcTypeLevelString $ mkString (runProperName ctorName)))
                           ctorTy
                 , CaseAlternative [ ConstructorBinder ss Libs.C_Constructor [matchProduct] ]
                                   (unguarded (foldl' App (Constructor ss (Qualified (ByModuleName mn) ctorName)) ctorArgs))
                 , CaseAlternative [ ConstructorBinder ss (Qualified (ByModuleName mn) ctorName) matchCtor ]
                                   (unguarded (App (Constructor ss Libs.C_Constructor) mkProduct))
                 )
  
      makeProduct
        :: [SourceType]
        -> m (SourceType, Binder, [Expr], [Binder], Expr)
      makeProduct [] =
        pure (srcTypeConstructor Libs.NoArguments, NullBinder, [], [], Constructor ss Libs.C_NoArguments)
      makeProduct args = do
        (tys, bs1, es1, bs2, es2) <- unzip5 <$> traverse makeArg args
        pure ( foldr1 (\f -> srcTypeApp (srcTypeApp (srcTypeConstructor Libs.Product) f)) tys
             , foldr1 (\b1 b2 -> ConstructorBinder ss Libs.C_Product [b1, b2]) bs1
             , es1
             , bs2
             , foldr1 (\e1 -> App (App (Constructor ss Libs.C_Product) e1)) es2
             )
  
      makeArg :: SourceType -> m (SourceType, Binder, Expr, Binder, Expr)
      makeArg arg = do
        argName <- freshIdent "arg"
        pure ( srcTypeApp (srcTypeConstructor Libs.Argument) arg
             , ConstructorBinder ss Libs.C_Argument [ VarBinder ss argName ]
             , Var ss (Qualified (BySourcePos $ spanStart ss) argName)
             , VarBinder ss argName
             , App (Constructor ss Libs.C_Argument) (Var ss (Qualified (BySourcePos $ spanStart ss) argName))
             )
  
      underBinder :: (Binder -> Binder) -> CaseAlternative -> CaseAlternative
      underBinder f (CaseAlternative bs e) = CaseAlternative (map f bs) e
  
      underExpr :: (Expr -> Expr) -> CaseAlternative -> CaseAlternative
      underExpr f (CaseAlternative b [MkUnguarded e]) = CaseAlternative b (unguarded (f e))
      underExpr _ _ = internalError "underExpr: expected unguarded alternative"
  
      toRepTy :: [SourceType] -> SourceType
      toRepTy [] = srcTypeConstructor Libs.NoConstructors
      toRepTy [only] = only
      toRepTy ctors = foldr1 (\f -> srcTypeApp (srcTypeApp (srcTypeConstructor Libs.Sum) f)) ctors
  
  checkIsWildcard :: MonadError MultipleErrors m => SourceSpan -> ProperName 'TypeName -> SourceType -> m ()
  checkIsWildcard _ _ (TypeWildcard _ UnnamedWildcard) = return ()
  checkIsWildcard ss tyConNm _ =
    throwError . errorMessage' ss $ ExpectedWildcard tyConNm
  
  deriveNewtype
    :: forall m
     . MonadError MultipleErrors m
    => Declaration
    -> [SourceType]
    -> m ([Declaration], SourceType)
  deriveNewtype tyCon tyConArgs =
    case tyCon of
      DataDeclaration (ss', _) Data name _ _ ->
        throwError . errorMessage' ss' $ CannotDeriveNewtypeForData name
      DataDeclaration _ Newtype name args dctors -> do
        (_, (_, ty)) <- checkNewtype name dctors
        let subst = zipWith ((,) . fst) args tyConArgs
        return ([], replaceAllTypeVars subst ty)
      _ -> internalError "deriveNewtype: expected DataDeclaration"
  
  findTypeDecl
    :: (MonadError MultipleErrors m)
    => SourceSpan
    -> ProperName 'TypeName
    -> [Declaration]
    -> m Declaration
  findTypeDecl ss tyConNm = note (errorMessage' ss $ CannotFindDerivingType tyConNm) . find isTypeDecl
    where
    isTypeDecl :: Declaration -> Bool
    isTypeDecl (DataDeclaration _ _ nm _ _) = nm == tyConNm
    isTypeDecl _ = False
  