-- | A minor cleanup pass that runs after defunctorisation and applies
-- any type abbreviations. After this, the program consists entirely
-- value bindings.
module Futhark.Internalise.ApplyTypeAbbrs (transformProg) where

import Control.Monad.Identity
import Data.Map.Strict qualified as M
import Data.Maybe (mapMaybe)
import Language.Futhark
import Language.Futhark.Semantic (TypeBinding (..))
import Language.Futhark.Traversals
import Language.Futhark.TypeChecker.Types

type Types = M.Map VName (Subst StructRetType)

getTypes :: Types -> [Dec] -> Types
getTypes types [] = types
getTypes types (TypeDec typebind : ds) = do
  let (TypeBind name l tparams _ (Info (RetType dims t)) _ _) = typebind
      tbinding = TypeAbbr l tparams $ RetType dims $ applySubst (`M.lookup` types) t
      types' = M.insert name (substFromAbbr tbinding) types
  getTypes types' ds
getTypes types (_ : ds) =
  getTypes types ds

-- Perform a given substitution on the types in a pattern.
substPat :: (t -> t) -> Pat t -> Pat t
substPat f pat = case pat of
  TuplePat pats loc -> TuplePat (map (substPat f) pats) loc
  RecordPat fs loc -> RecordPat (map substField fs) loc
    where
      substField (n, p) = (n, substPat f p)
  PatParens p loc -> PatParens (substPat f p) loc
  PatAttr attr p loc -> PatAttr attr (substPat f p) loc
  Id vn (Info tp) loc -> Id vn (Info $ f tp) loc
  Wildcard (Info tp) loc -> Wildcard (Info $ f tp) loc
  PatAscription p _ _ -> substPat f p
  PatLit e (Info tp) loc -> PatLit e (Info $ f tp) loc
  PatConstr n (Info tp) ps loc -> PatConstr n (Info $ f tp) ps loc

removeTypeVariablesInType :: Types -> StructType -> StructType
removeTypeVariablesInType types =
  applySubst (`M.lookup` types)

substEntry :: Types -> EntryPoint -> EntryPoint
substEntry types (EntryPoint params ret) =
  EntryPoint (map onEntryParam params) (onEntryType ret)
  where
    onEntryParam (EntryParam v t) =
      EntryParam v $ onEntryType t
    onEntryType (EntryType t te) =
      EntryType (removeTypeVariablesInType types t) te

-- Remove all type variables and type abbreviations from a value binding.
removeTypeVariables :: Types -> ValBind -> ValBind
removeTypeVariables types valbind = do
  let (ValBind entry _ _ (Info (RetType dims rettype)) _ pats body _ _ _) = valbind
      mapper =
        ASTMapper
          { mapOnExp = onExp,
            mapOnName = pure,
            mapOnStructType = pure . applySubst (`M.lookup` types),
            mapOnParamType = pure . applySubst (`M.lookup` types),
            mapOnResRetType = pure . applySubst (`M.lookup` types)
          }
      onExp = astMap mapper

  let body' = runIdentity $ onExp body

  valbind
    { valBindRetType = Info (applySubst (`M.lookup` types) $ RetType dims rettype),
      valBindParams = map (substPat $ applySubst (`M.lookup` types)) pats,
      valBindEntryPoint = fmap (substEntry types) <$> entry,
      valBindBody = body'
    }

-- | Apply type abbreviations from a list of top-level declarations. A
-- module-free input program is expected, so only value declarations
-- and type declaration are accepted.
transformProg :: (Monad m) => [Dec] -> m [ValBind]
transformProg decs =
  let types = getTypes mempty decs
      onDec (ValDec valbind) = Just $ removeTypeVariables types valbind
      onDec _ = Nothing
   in pure $ mapMaybe onDec decs