TypedDesugar.hs

{- |
Module      : Language.Egison.Type.TypedDesugar
Licence     : MIT

This module implements Phase 7 of the processing flow: TypedDesugar.
It orchestrates type-driven transformations on TIExpr (Typed Internal Expressions)
by calling specialized expansion modules.

Type-Driven Transformations (Phase 7):
  1. Type class dictionary passing (via TypeClassExpand)
     - Instance selection based on types
     - Method call concretization
  2. Type information optimization and embedding
     - Preserve type info for better error messages during evaluation
     - Each node in TIExpr contains its type

Type information is preserved throughout desugaring, enabling:
  - Better runtime error messages with type information
  - Type-based dispatch during evaluation
  - Debugging support with type annotations
-}

module Language.Egison.Type.TypedDesugar
  ( desugarTypedExprT
  , desugarTypedTopExprT
  , desugarTypedTopExprT_TensorMapOnly
  , desugarTypedTopExprT_TypeClassOnly
  ) where

import           Language.Egison.Data       (EvalM)
import           Language.Egison.EvalState  (MonadEval(..))
import           Language.Egison.IExpr      (TIExpr(..), TIExprNode(..), TITopExpr(..), extractNameFromVar, stringToVar)
import           Language.Egison.Type.Env   (lookupEnv)
import           Language.Egison.Type.TensorMapInsertion (insertTensorMaps)
import           Language.Egison.Type.Types (Type(..), TypeScheme(..))
import           Language.Egison.Type.TypeClassExpand (expandTypeClassMethodsT, expandTypeClassMethodsInPattern, addDictionaryParametersT, applyConcreteConstraintDictionaries, applyConcreteConstraintDictionariesInPattern, fixUnboundDictRefs)

-- | Wrap a TIExpr with TIReshape when the type scheme demands a concrete
-- CAS scalar type (Integer, Frac _, Poly _ _, Term _ _, Factor). Skip for
-- polymorphic schemes (any type variables or class constraints), non-CAS
-- types, and TMathValue (which is the most general — reshape is a no-op).
-- This is the post-typecheck elaboration step; placing it after type class
-- expansion preserves inner-method dispatch context.
maybeReshape :: TypeScheme -> TIExpr -> TIExpr
maybeReshape sch@(Forall vars constraints ty) tiexpr
  | not (null vars) || not (null constraints) = tiexpr
  | isReshapeTarget ty = TIExpr sch (TIReshape ty tiexpr)
  | otherwise = tiexpr
  where
    isReshapeTarget :: Type -> Bool
    isReshapeTarget TInt        = True
    isReshapeTarget TFactor     = True
    isReshapeTarget (TFrac _)   = True
    isReshapeTarget (TPoly _ _) = True
    isReshapeTarget (TTerm _ _) = True
    isReshapeTarget _           = False

-- | Desugar a typed expression (TIExpr) with type-driven transformations
-- This function orchestrates the transformation pipeline:
--   1. Insert tensorMap where needed (TensorMapInsertion)
--   2. Expand type class methods (dictionary passing)
--
-- The order matters: tensorMap insertion should happen before type class expansion
-- because after tensorMap insertion, argument types (scalar vs tensor) are determined,
-- which allows type class expansion to use unifyStrict for instance selection.
desugarTypedExprT :: TIExpr -> EvalM TIExpr
desugarTypedExprT tiexpr = do
  -- Step 1: Insert tensorMap where needed
  tiexpr' <- insertTensorMaps tiexpr

  -- Step 2: Expand type class methods (dictionary passing)
  tiexpr'' <- expandTypeClassMethodsT tiexpr'

  return tiexpr''

-- | Desugar a top-level typed expression (TITopExpr)
-- This is the main entry point for Phase 7 transformations.
desugarTypedTopExprT :: TITopExpr -> EvalM (Maybe TITopExpr)
desugarTypedTopExprT topExpr = case topExpr of
  TIDefine scheme var tiexpr -> do
    tiexpr' <- desugarTypedExprT tiexpr
    -- Apply dictionaries to right-hand side if it has concrete type constraints
    tiexpr'' <- applyConcreteConstraintDictionaries tiexpr'
    -- Add dictionary parameters for constrained functions
    tiexpr''' <- addDictionaryParametersT scheme tiexpr''
    -- Repair any dictionary access left unbound (fall back to runtime dispatch)
    classEnv <- getClassEnv
    let tiexprFixed = fixUnboundDictRefs classEnv tiexpr'''
    -- Insert TIReshape from type annotation (post-typecheck elaboration)
    let tiexprFinal = maybeReshape scheme tiexprFixed
    return $ Just (TIDefine scheme var tiexprFinal)

  TITest tiexpr -> do
    tiexpr' <- desugarTypedExprT tiexpr
    classEnv <- getClassEnv
    return $ Just (TITest (fixUnboundDictRefs classEnv tiexpr'))

  TIExecute tiexpr -> do
    tiexpr' <- desugarTypedExprT tiexpr
    classEnv <- getClassEnv
    return $ Just (TIExecute (fixUnboundDictRefs classEnv tiexpr'))
  
  TILoadFile path -> 
    return $ Just (TILoadFile path)
  
  TILoad lib -> 
    return $ Just (TILoad lib)
  
  TIDefineMany bindings -> do
    bindings' <- mapM (\(var, tiexpr) -> do
      tiexpr' <- desugarTypedExprT tiexpr
      -- Add dictionary parameters using the variable's type scheme from TypeEnv
      -- This is important for dictionary definitions where the expression (hash)
      -- may not have constraints, but the variable has constraints in its type scheme
      typeEnv <- getTypeEnv
      let varName = extractNameFromVar var
          scheme = case lookupEnv (stringToVar varName) typeEnv of
                     Just ts -> ts  -- Use type scheme from environment
                     Nothing -> tiScheme tiexpr'  -- Fallback to expression's scheme
      tiexpr'' <- addDictionaryParametersT scheme tiexpr'
      classEnv <- getClassEnv
      return (var, fixUnboundDictRefs classEnv tiexpr'')) bindings
    return $ Just (TIDefineMany bindings')
  
  TIDeclareSymbol names ty ->
    -- Symbol declarations don't need type-driven transformations
    return $ Just (TIDeclareSymbol names ty)
  
  TIPatternFunctionDecl name typeScheme params retType body -> do
    -- Pattern function declarations: apply type class expansion and dictionary application to body
    body' <- expandTypeClassMethodsInPattern body
    body'' <- applyConcreteConstraintDictionariesInPattern body'
    return $ Just (TIPatternFunctionDecl name typeScheme params retType body'')

-- | Desugar a top-level typed expression with TensorMap insertion only
-- This is used for --dump-ti (intermediate dump after TensorMap insertion)
desugarTypedTopExprT_TensorMapOnly :: TITopExpr -> EvalM (Maybe TITopExpr)
desugarTypedTopExprT_TensorMapOnly topExpr = case topExpr of
  TIDefine scheme var tiexpr -> do
    -- Only insert tensorMap (no type class expansion)
    tiexpr' <- insertTensorMaps tiexpr
    return $ Just (TIDefine scheme var tiexpr')

  TITest tiexpr -> do
    tiexpr' <- insertTensorMaps tiexpr
    return $ Just (TITest tiexpr')

  TIExecute tiexpr -> do
    tiexpr' <- insertTensorMaps tiexpr
    return $ Just (TIExecute tiexpr')

  TILoadFile path ->
    return $ Just (TILoadFile path)

  TILoad lib ->
    return $ Just (TILoad lib)

  TIDefineMany bindings -> do
    bindings' <- mapM (\(var, tiexpr) -> do
      tiexpr' <- insertTensorMaps tiexpr
      return (var, tiexpr')) bindings
    return $ Just (TIDefineMany bindings')

  TIDeclareSymbol names ty ->
    return $ Just (TIDeclareSymbol names ty)
  
  TIPatternFunctionDecl name typeScheme params retType body ->
    -- Pattern function declarations: TensorMap insertion only
    return $ Just (TIPatternFunctionDecl name typeScheme params retType body)

-- | Expand type class methods only (assumes TensorMap insertion is already done)
-- This is used internally to perform type class expansion after TensorMap insertion
desugarTypedTopExprT_TypeClassOnly :: TITopExpr -> EvalM (Maybe TITopExpr)
desugarTypedTopExprT_TypeClassOnly topExpr = case topExpr of
  TIDefine scheme var tiexpr -> do
    -- Only expand type class methods (assumes tensorMap is already inserted)
    tiexpr' <- expandTypeClassMethodsT tiexpr
    -- Apply dictionaries to right-hand side if it has concrete type constraints
    tiexpr'' <- applyConcreteConstraintDictionaries tiexpr'
    -- Add dictionary parameters for constrained functions
    tiexpr''' <- addDictionaryParametersT scheme tiexpr''
    -- Repair any dictionary access left unbound (fall back to runtime dispatch)
    classEnv <- getClassEnv
    let tiexprFixed = fixUnboundDictRefs classEnv tiexpr'''
    -- Insert TIReshape from type annotation (post-typecheck elaboration)
    let tiexprFinal = maybeReshape scheme tiexprFixed
    return $ Just (TIDefine scheme var tiexprFinal)

  TITest tiexpr -> do
    tiexpr' <- expandTypeClassMethodsT tiexpr
    tiexpr'' <- applyConcreteConstraintDictionaries tiexpr'
    classEnv <- getClassEnv
    return $ Just (TITest (fixUnboundDictRefs classEnv tiexpr''))

  TIExecute tiexpr -> do
    tiexpr' <- expandTypeClassMethodsT tiexpr
    tiexpr'' <- applyConcreteConstraintDictionaries tiexpr'
    classEnv <- getClassEnv
    return $ Just (TIExecute (fixUnboundDictRefs classEnv tiexpr''))

  TILoadFile path ->
    return $ Just (TILoadFile path)

  TILoad lib ->
    return $ Just (TILoad lib)

  TIDefineMany bindings -> do
    bindings' <- mapM (\(var, tiexpr) -> do
      tiexpr' <- expandTypeClassMethodsT tiexpr
      -- Add dictionary parameters using the variable's type scheme from TypeEnv
      typeEnv <- getTypeEnv
      let varName = extractNameFromVar var
          scheme = case lookupEnv (stringToVar varName) typeEnv of
                     Just ts -> ts
                     Nothing -> tiScheme tiexpr'
      tiexpr'' <- addDictionaryParametersT scheme tiexpr'
      classEnv <- getClassEnv
      return (var, fixUnboundDictRefs classEnv tiexpr'')) bindings
    return $ Just (TIDefineMany bindings')
  
  TIDeclareSymbol names ty ->
    return $ Just (TIDeclareSymbol names ty)
  
  TIPatternFunctionDecl name typeScheme params retType body -> do
    -- Pattern function declarations: expand type class methods and apply dictionaries in body
    body' <- expandTypeClassMethodsInPattern body
    body'' <- applyConcreteConstraintDictionariesInPattern body'
    return $ Just (TIPatternFunctionDecl name typeScheme params retType body'')