packages feed

egison-5.0.0: hs-src/Language/Egison/Type/TypeClassExpand.hs

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

This module expands type class method calls using type information from TIExpr.
It transforms TIExpr to TIExpr, replacing type class method calls with
dictionary-based dispatch.

Pipeline: Phase 8 (TypedDesugar) - TypeClassExpand (first step)
This is executed before TensorMapInsertion to resolve type class methods
to concrete functions first.

For example, if we have:
  class Eq a where (==) : a -> a -> Bool
  instance Eq Integer where (==) x y := x = y

Then a call like:
  autoEq 1 2  (with type constraint: Eq Integer)
becomes:
  eqIntegerEq 1 2  (dictionary-based dispatch)

This eliminates the need for runtime dispatch functions like resolveEq.
-}

module Language.Egison.Type.TypeClassExpand
  ( expandTypeClassMethodsT
  , expandTypeClassMethodsInPattern
  , addDictionaryParametersT
  , applyConcreteConstraintDictionaries
  , applyConcreteConstraintDictionariesInPattern
  ) where

import           Data.Char                  (toLower)
import           Data.List                  (find)
import           Data.Maybe                 (mapMaybe)
import           Data.Text                  (pack)
import           Control.Monad              (mplus)
import qualified Data.Set                   as Set

import           Language.Egison.AST        (ConstantExpr(..))
import           Language.Egison.Data       (EvalM)
import           Language.Egison.EvalState  (MonadEval(..))
import           Language.Egison.IExpr      (TIExpr(..), TIExprNode(..), IExpr(..), stringToVar,
                                             Index(..), tiExprType, tiScheme, tiExprNode,
                                             TIPattern(..), TIPatternNode(..), TILoopRange(..))
import           Language.Egison.Type.Env  (ClassEnv(..), ClassInfo(..), InstanceInfo(..),
                                             lookupInstances, lookupClass, lookupEnv)
import qualified Language.Egison.Type.Types as Types
import           Language.Egison.Type.Types (Type(..), TyVar(..), TypeScheme(..), Constraint(..), typeToName, typeConstructorName,
                                            sanitizeMethodName, freeTyVars)
import           Language.Egison.Type.Instance (findMatchingInstanceForType)

-- ============================================================================
-- Helper Functions (shared across the module)
-- ============================================================================

-- | Extract type variable substitutions from instance type and actual type
-- Example: [a] -> [[Integer]] gives [(a, [Integer])]
extractTypeSubstitutions :: Type -> Type -> [(TyVar, Type)]
extractTypeSubstitutions instTy actualTy = go instTy actualTy
  where
    go (TVar v) actual = [(v, actual)]
    go (TCollection instElem) (TCollection actualElem) = go instElem actualElem
    go (TTuple instTypes) (TTuple actualTypes)
      | length instTypes == length actualTypes =
          concatMap (\(i, a) -> go i a) (zip instTypes actualTypes)
    go (TInductive _ instArgs) (TInductive _ actualArgs)
      | length instArgs == length actualArgs =
          concatMap (\(i, a) -> go i a) (zip instArgs actualArgs)
    go (TTensor instElem) (TTensor actualElem) = go instElem actualElem
    go (TFun instArg instRet) (TFun actualArg actualRet) =
      go instArg actualArg ++ go instRet actualRet
    go (THash instK instV) (THash actualK actualV) =
      go instK actualK ++ go instV actualV
    go (TMatcher instT) (TMatcher actualT) = go instT actualT
    go (TIO instT) (TIO actualT) = go instT actualT
    go (TIORef instT) (TIORef actualT) = go instT actualT
    go TPort TPort = []
    go _ _ = []

-- | Apply type substitutions to a constraint
applySubstsToConstraint :: [(TyVar, Type)] -> Constraint -> Constraint
applySubstsToConstraint substs (Constraint cName cType) =
  Constraint cName (applySubstsToType substs cType)

-- | Apply type substitutions to a type
applySubstsToType :: [(TyVar, Type)] -> Type -> Type
applySubstsToType substs = go
  where
    go t@(TVar v) = case lookup v substs of
                      Just newType -> newType
                      Nothing -> t
    go TInt = TInt
    go TFloat = TFloat
    go TBool = TBool
    go TChar = TChar
    go TString = TString
    go (TCollection t) = TCollection (go t)
    go (TTuple ts) = TTuple (map go ts)
    go (TInductive name ts) = TInductive name (map go ts)
    go (TTensor t) = TTensor (go t)
    go (THash k v) = THash (go k) (go v)
    go (TMatcher t) = TMatcher (go t)
    go (TFun t1 t2) = TFun (go t1) (go t2)
    go (TIO t) = TIO (go t)
    go (TIORef t) = TIORef (go t)
    go TPort = TPort
    go TAny = TAny

-- | Get the arity of a function type (number of parameters)
getMethodArity :: Type -> Int
getMethodArity (TFun _ t2) = 1 + getMethodArity t2
getMethodArity _ = 0

-- | Get parameter types from a function type
getParamTypes :: Type -> [Type]
getParamTypes (TFun t1 t2) = t1 : getParamTypes t2
getParamTypes _ = []

-- | Apply N parameters to a function type and get the result type
-- applyParamsToType (a -> b -> c) 2 = c
-- applyParamsToType (a -> b -> c) 1 = b -> c
applyParamsToType :: Type -> Int -> Type
applyParamsToType (TFun _ t2) n
  | n > 0 = applyParamsToType t2 (n - 1)
applyParamsToType t _ = t  -- n == 0 or no more function types

-- | Lowercase first character of a string
lowerFirst :: String -> String
lowerFirst [] = []
lowerFirst (c:cs) = toLower c : cs

-- | Find a constraint that provides the given method
findConstraintForMethod :: ClassEnv -> String -> [Constraint] -> Maybe Constraint
findConstraintForMethod env methodName cs =
  find (\(Constraint className _) ->
    case lookupClass className env of
      Just classInfo -> methodName `elem` map fst (classMethods classInfo)
      Nothing -> False
  ) cs

-- ============================================================================
-- Main Type Class Expansion
-- ============================================================================

-- | Expand type class method calls in a typed expression (TIExpr)
-- This function recursively processes TIExpr and replaces type class method calls
-- with dictionary-based dispatch.
expandTypeClassMethodsT :: TIExpr -> EvalM TIExpr
expandTypeClassMethodsT tiExpr = do
  classEnv <- getClassEnv
  let scheme = tiScheme tiExpr
  -- Recursively process the TIExprNode with constraint information
  expandedNode <- expandTIExprNodeWithConstraints classEnv scheme (tiExprNode tiExpr)
  return $ TIExpr scheme expandedNode
  where
    -- Expand TIExprNode with constraint information from TypeScheme
    -- Note: Constraints from parent are not propagated - each node uses its own constraints
    expandTIExprNodeWithConstraints :: ClassEnv -> TypeScheme -> TIExprNode -> EvalM TIExprNode
    expandTIExprNodeWithConstraints classEnv' (Forall _vars _constraints _ty) node =
      expandTIExprNode classEnv' node

    -- Expand TIExprNode without parent constraints
    -- Each child expression uses only its own constraints from type inference
    expandTIExprNode :: ClassEnv -> TIExprNode -> EvalM TIExprNode
    expandTIExprNode classEnv' node = case node of
      -- Constants and variables: no expansion needed at node level
      -- (TIVarExpr expansion is handled at TIExpr level in expandTIExprWithConstraints)
      TIConstantExpr c -> return $ TIConstantExpr c
      TIVarExpr name -> return $ TIVarExpr name
      
      -- Lambda expressions: process body with its own constraints only
      TILambdaExpr mVar params body -> do
        -- Use only the body's own constraints (no parent constraints)
        -- Type inference has already assigned correct constraints to each expression
        body' <- expandTIExprWithConstraints classEnv' body
        return $ TILambdaExpr mVar params body'
      
      -- Application: check if it's a method call or constrained function call
      TIApplyExpr func args -> do
        -- First, expand the arguments (each uses its own constraints)
        args' <- mapM (expandTIExprWithConstraints classEnv') args

        case tiExprNode func of
          TIVarExpr methodName -> do
            -- Try to resolve if func is a method call using func's own constraints
            let (Forall _ funcConstraints _) = tiScheme func
            resolved <- tryResolveMethodCall classEnv' funcConstraints methodName args'
            case resolved of
              Just result -> return result
              Nothing -> do
                -- Not a method call - process recursively
                -- Note: Dictionary application for constrained functions
                -- is handled in TIVarExpr case of expandTIExprWithConstraints
                func' <- expandTIExprWithConstraints classEnv' func
                return $ TIApplyExpr func' args'
          _ -> do
            -- Not a simple variable: process recursively
            func' <- expandTIExprWithConstraints classEnv' func
            return $ TIApplyExpr func' args'
      
      -- Collections
      TITupleExpr exprs -> do
        exprs' <- mapM (expandTIExprWithConstraints classEnv') exprs
        return $ TITupleExpr exprs'

      TICollectionExpr exprs -> do
        exprs' <- mapM (expandTIExprWithConstraints classEnv') exprs
        return $ TICollectionExpr exprs'

      -- Control flow
      TIIfExpr cond thenExpr elseExpr -> do
        cond' <- expandTIExprWithConstraints classEnv' cond
        thenExpr' <- expandTIExprWithConstraints classEnv' thenExpr
        elseExpr' <- expandTIExprWithConstraints classEnv' elseExpr
        return $ TIIfExpr cond' thenExpr' elseExpr'

      -- Let bindings
      TILetExpr bindings body -> do
        bindings' <- mapM (\(v, e) -> do
          e' <- expandTIExprWithConstraints classEnv' e
          return (v, e')) bindings
        body' <- expandTIExprWithConstraints classEnv' body
        return $ TILetExpr bindings' body'

      TILetRecExpr bindings body -> do
        bindings' <- mapM (\(v, e) -> do
          e' <- expandTIExprWithConstraints classEnv' e
          return (v, e')) bindings
        body' <- expandTIExprWithConstraints classEnv' body
        return $ TILetRecExpr bindings' body'

      TISeqExpr e1 e2 -> do
        e1' <- expandTIExprWithConstraints classEnv' e1
        e2' <- expandTIExprWithConstraints classEnv' e2
        return $ TISeqExpr e1' e2'

      -- Collections
      TIConsExpr h t -> do
        h' <- expandTIExprWithConstraints classEnv' h
        t' <- expandTIExprWithConstraints classEnv' t
        return $ TIConsExpr h' t'

      TIJoinExpr l r -> do
        l' <- expandTIExprWithConstraints classEnv' l
        r' <- expandTIExprWithConstraints classEnv' r
        return $ TIJoinExpr l' r'
      
      TIHashExpr pairs -> do
        -- Dictionary hashes: process keys but NOT values
        -- Values should remain as simple method references
        pairs' <- mapM (\(k, v) -> do
          k' <- expandTIExprWithConstraints classEnv' k
          -- Do NOT process v - dictionary values should not be expanded
          return (k', v)) pairs
        return $ TIHashExpr pairs'

      TIVectorExpr exprs -> do
        exprs' <- mapM (expandTIExprWithConstraints classEnv') exprs
        return $ TIVectorExpr exprs'

      -- More lambda-like constructs
      TIMemoizedLambdaExpr vars body -> do
        body' <- expandTIExprWithConstraints classEnv' body
        return $ TIMemoizedLambdaExpr vars body'

      TICambdaExpr var body -> do
        body' <- expandTIExprWithConstraints classEnv' body
        return $ TICambdaExpr var body'

      TIWithSymbolsExpr syms body -> do
        body' <- expandTIExprWithConstraints classEnv' body
        return $ TIWithSymbolsExpr syms body'

      TIDoExpr bindings body -> do
        bindings' <- mapM (\(v, e) -> do
          e' <- expandTIExprWithConstraints classEnv' e
          return (v, e')) bindings
        body' <- expandTIExprWithConstraints classEnv' body
        return $ TIDoExpr bindings' body'

      -- Pattern matching
      TIMatchExpr mode target matcher clauses -> do
        target' <- expandTIExprWithConstraints classEnv' target
        matcher' <- expandTIExprWithConstraints classEnv' matcher
        clauses' <- mapM (\(pat, body) -> do
          pat' <- expandTIPattern classEnv' pat
          body' <- expandTIExprWithConstraints classEnv' body
          return (pat', body')) clauses
        return $ TIMatchExpr mode target' matcher' clauses'

      TIMatchAllExpr mode target matcher clauses -> do
        target' <- expandTIExprWithConstraints classEnv' target
        matcher' <- expandTIExprWithConstraints classEnv' matcher
        clauses' <- mapM (\(pat, body) -> do
          pat' <- expandTIPattern classEnv' pat
          body' <- expandTIExprWithConstraints classEnv' body
          return (pat', body')) clauses
        return $ TIMatchAllExpr mode target' matcher' clauses'

      -- Tensor operations
      TITensorMapExpr func tensor -> do
        func' <- expandTIExprWithConstraints classEnv' func
        tensor' <- expandTIExprWithConstraints classEnv' tensor
        return $ TITensorMapExpr func' tensor'

      TITensorMap2Expr func t1 t2 -> do
        func' <- expandTIExprWithConstraints classEnv' func
        t1' <- expandTIExprWithConstraints classEnv' t1
        t2' <- expandTIExprWithConstraints classEnv' t2
        return $ TITensorMap2Expr func' t1' t2'

      TITensorMap2WedgeExpr func t1 t2 -> do
        func' <- expandTIExprWithConstraints classEnv' func
        t1' <- expandTIExprWithConstraints classEnv' t1
        t2' <- expandTIExprWithConstraints classEnv' t2
        return $ TITensorMap2WedgeExpr func' t1' t2'

      TIGenerateTensorExpr func shape -> do
        func' <- expandTIExprWithConstraints classEnv' func
        shape' <- expandTIExprWithConstraints classEnv' shape
        return $ TIGenerateTensorExpr func' shape'

      TITensorExpr shape elems -> do
        shape' <- expandTIExprWithConstraints classEnv' shape
        elems' <- expandTIExprWithConstraints classEnv' elems
        return $ TITensorExpr shape' elems'

      TITensorContractExpr tensor -> do
        tensor' <- expandTIExprWithConstraints classEnv' tensor
        return $ TITensorContractExpr tensor'

      TITransposeExpr perm tensor -> do
        perm' <- expandTIExprWithConstraints classEnv' perm
        tensor' <- expandTIExprWithConstraints classEnv' tensor
        return $ TITransposeExpr perm' tensor'

      TIFlipIndicesExpr tensor -> do
        tensor' <- expandTIExprWithConstraints classEnv' tensor
        return $ TIFlipIndicesExpr tensor'

      -- Quote expressions
      TIQuoteExpr e -> do
        e' <- expandTIExprWithConstraints classEnv' e
        return $ TIQuoteExpr e'

      TIQuoteSymbolExpr e -> do
        e' <- expandTIExprWithConstraints classEnv' e
        return $ TIQuoteSymbolExpr e'

      -- Indexed expressions
      TISubrefsExpr b base ref -> do
        base' <- expandTIExprWithConstraints classEnv' base
        ref' <- expandTIExprWithConstraints classEnv' ref
        return $ TISubrefsExpr b base' ref'
      
      TISuprefsExpr b base ref -> do
        base' <- expandTIExprWithConstraints classEnv' base
        ref' <- expandTIExprWithConstraints classEnv' ref
        return $ TISuprefsExpr b base' ref'

      TIUserrefsExpr b base ref -> do
        base' <- expandTIExprWithConstraints classEnv' base
        ref' <- expandTIExprWithConstraints classEnv' ref
        return $ TIUserrefsExpr b base' ref'

      -- Other cases: return unchanged for now
      TIInductiveDataExpr name exprs -> do
        exprs' <- mapM (expandTIExprWithConstraints classEnv') exprs
        return $ TIInductiveDataExpr name exprs'

      TIMatcherExpr patDefs -> do
        -- Expand expressions inside matcher definitions
        -- patDefs is a list of (PrimitivePatPattern, TIExpr, [TIBindingExpr])
        -- where TIBindingExpr is (IPrimitiveDataPattern, TIExpr)
        patDefs' <- mapM (\(pat, matcherExpr, bindings) -> do
          -- Expand the next-matcher expression
          matcherExpr' <- expandTIExprWithConstraints classEnv' matcherExpr
          -- Expand expressions in primitive-data-match clauses
          bindings' <- mapM (\(dp, expr) -> do
            expr' <- expandTIExprWithConstraints classEnv' expr
            return (dp, expr')) bindings
          return (pat, matcherExpr', bindings')) patDefs
        return $ TIMatcherExpr patDefs'
      TIIndexedExpr override base indices -> do
        base' <- expandTIExprWithConstraints classEnv' base
        -- Expand indices (which are already typed as TIExpr)
        indices' <- mapM (traverse (\tiexpr -> expandTIExprWithConstraints classEnv' tiexpr)) indices
        return $ TIIndexedExpr override base' indices'

      TIWedgeApplyExpr func args -> do
        func' <- expandTIExprWithConstraints classEnv' func
        args' <- mapM (expandTIExprWithConstraints classEnv') args
        return $ TIWedgeApplyExpr func' args'
      
      TIFunctionExpr names -> return $ TIFunctionExpr names  -- Built-in function, no expansion needed
    
    -- Helper: expand a TIExpr using only its own constraints
    -- Parent constraints are not passed to avoid constraint accumulation
    expandTIExprWithConstraints :: ClassEnv -> TIExpr -> EvalM TIExpr
    expandTIExprWithConstraints classEnv' expr = do
      let scheme@(Forall _ exprConstraints exprType) = tiScheme expr
          -- Use only the expression's own constraints
          -- Type inference has already assigned correct constraints to each expression
          allConstraints = exprConstraints

      -- Special handling for TIVarExpr: eta-expand methods or apply dictionaries
      expandedNode <- case tiExprNode expr of
        TIVarExpr varName -> do
          -- Check if this is a type class method
          case findConstraintForMethod classEnv' varName allConstraints of
            Just (Constraint className tyArg) -> do
              -- Get method type to determine arity
              typeEnv <- getTypeEnv
              case lookupEnv (stringToVar varName) typeEnv of
                Just (Forall _ _ _ty) -> do
                  -- Use the expression's actual type (exprType) instead of the method's declared type (ty)
                  -- because eta-expansion should create parameters matching the expected usage context
                  let arity = getMethodArity exprType
                      paramTypes = getParamTypes exprType
                      paramNames = ["etaVar" ++ show i | i <- [1..arity]]
                      paramVars = map stringToVar paramNames
                      paramExprs = zipWith (\n t -> TIExpr (Forall [] [] t) (TIVarExpr n)) paramNames paramTypes
                      methodKey = sanitizeMethodName varName
                  
                  -- Determine dictionary name based on type
                  case tyArg of
                    TVar (TyVar _v) -> do
                      -- Type variable: use dictionary parameter name (without type parameter)
                      typeEnv <- getTypeEnv
                      let dictParamName = "dict_" ++ className
                      -- Look up dictionary type from type environment
                      dictHashType <- case lookupEnv (stringToVar dictParamName) typeEnv of
                        Just (Forall _ _ dictType) -> return dictType
                        Nothing -> return $ THash TString TAny  -- Fallback
                      -- Get method type from ClassEnv instead of dictHashType
                      let methodType = getMethodTypeFromClass classEnv' className methodKey tyArg
                          methodConstraint = Constraint className tyArg
                          methodScheme = Forall (Set.toList $ freeTyVars tyArg) [methodConstraint] methodType
                          dictExpr = TIExpr (Forall [] [] dictHashType) (TIVarExpr dictParamName)
                          indexExpr = TIExpr (Forall [] [] TString)
                                            (TIConstantExpr (StringExpr (pack methodKey)))
                          dictAccess = TIExpr methodScheme $
                                       TIIndexedExpr False dictExpr [Sub indexExpr]
                          -- Calculate result type after applying all parameters
                          resultType = applyParamsToType methodType (length paramExprs)
                          -- Fully applied results don't need constraints
                          bodyScheme = case resultType of
                                         TFun _ _ -> methodScheme  -- Partial application
                                         _ -> Forall [] [] resultType  -- Fully applied: no constraints
                          body = TIExpr bodyScheme (TIApplyExpr dictAccess paramExprs)
                      return $ TILambdaExpr Nothing paramVars body
                    _ -> do
                      -- Concrete type: find matching instance
                      let instances = lookupInstances className classEnv'
                      case findMatchingInstanceForType tyArg instances of
                        Just inst -> do
                          -- Found instance: eta-expand with concrete dictionary
                          typeEnv <- getTypeEnv
                          let instTypeName = typeConstructorName (instType inst)
                              dictName = lowerFirst className ++ instTypeName

                          -- Look up dictionary type from type environment
                          dictHashType <- case lookupEnv (stringToVar dictName) typeEnv of
                            Just (Forall _ _ dictType) -> return dictType
                            Nothing -> return $ THash TString TAny  -- Fallback

                          -- Get method type from ClassEnv instead of dictHashType
                          let methodType = getMethodTypeFromClass classEnv' className methodKey tyArg
                              methodConstraint = Constraint className tyArg
                              methodScheme = Forall (Set.toList $ freeTyVars tyArg) [methodConstraint] methodType

                          -- Check if instance has nested constraints
                          dictExprBase <- if null (instContext inst)
                            then do
                              -- No constraints: dictionary is a simple hash
                              return $ TIExpr (Forall [] [] dictHashType) (TIVarExpr dictName)
                            else do
                              -- Has constraints: dictionary is a function that returns a hash
                              -- Get the result type (should be the hash type after applying arguments)
                              let dictFuncType = case dictHashType of
                                    TFun _ resultType -> TFun dictHashType resultType
                                    _ -> TFun (THash TString TAny) dictHashType
                                  dictFuncExpr = TIExpr (Forall [] [] dictFuncType) (TIVarExpr dictName)
                              dictArgs <- mapM (resolveDictionaryArg classEnv') (instContext inst)
                              return $ TIExpr (Forall [] [] dictHashType) (TIApplyExpr dictFuncExpr dictArgs)

                          let indexExpr = TIExpr (Forall [] [] TString)
                                               (TIConstantExpr (StringExpr (pack methodKey)))
                              dictAccess = TIExpr methodScheme $
                                           TIIndexedExpr False dictExprBase [Sub indexExpr]
                              -- Calculate result type after applying all parameters
                              resultType = applyParamsToType methodType (length paramExprs)
                              -- Fully applied results don't need constraints
                              bodyScheme = case resultType of
                                             TFun _ _ -> methodScheme  -- Partial application
                                             _ -> Forall [] [] resultType  -- Fully applied: no constraints
                              body = TIExpr bodyScheme (TIApplyExpr dictAccess paramExprs)
                          return $ TILambdaExpr Nothing paramVars body
                        Nothing -> checkConstrainedVariable
                Nothing -> checkConstrainedVariable
            Nothing -> checkConstrainedVariable
          where
            -- Check if this is a constrained variable (not a method)
            -- IMPORTANT: Only apply dictionaries if the variable was DEFINED with constraints,
            -- not just if the expression has propagated constraints from usage context.
            checkConstrainedVariable = do
              typeEnv <- getTypeEnv
              -- Look up the variable's original type scheme from TypeEnv
              case lookupEnv (stringToVar varName) typeEnv of
                Just (Forall _ originalConstraints _)
                  | not (null originalConstraints) -> do
                      -- Variable was defined with constraints - apply dictionaries
                      -- Check if all constraints are on concrete types
                      let hasOnlyConcreteConstraints = all isConcreteConstraint exprConstraints
                      if hasOnlyConcreteConstraints
                        then do
                          -- This is a constrained variable with concrete types - apply dictionaries
                          dictArgs <- mapM (resolveDictionaryArg classEnv') exprConstraints
                          -- Create application: varName dict1 dict2 ...
                          let varExpr = TIExpr scheme (TIVarExpr varName)
                          return $ TIApplyExpr varExpr dictArgs
                        else do
                          -- Has type variable constraints - pass dictionary parameters
                          -- This handles recursive calls in polymorphic functions
                          -- Generate dictionary argument expressions for each constraint
                          let makeDict c =
                                let dictName = constraintToDictParam c
                                    dictType = TVar (TyVar "dict")
                                in TIExpr (Forall [] [] dictType) (TIVarExpr dictName)
                              dictArgs = map makeDict exprConstraints
                              varExpr = TIExpr scheme (TIVarExpr varName)
                          return $ TIApplyExpr varExpr dictArgs
                _ ->
                  -- Variable was defined without constraints, or not found in TypeEnv
                  -- Don't apply dictionaries - just process normally
                  expandTIExprNode classEnv' (tiExprNode expr)

            isConcreteConstraint (Constraint _ (TVar _)) = False
            isConcreteConstraint _ = True
        _ -> expandTIExprNode classEnv' (tiExprNode expr)

      return $ TIExpr scheme expandedNode

    -- Expand type class methods in patterns (no parent constraints)
    expandTIPattern :: ClassEnv -> TIPattern -> EvalM TIPattern
    expandTIPattern classEnv' (TIPattern scheme node) = do
      node' <- expandTIPatternNode classEnv' node
      return $ TIPattern scheme node'

    -- Expand pattern nodes recursively (no parent constraints)
    expandTIPatternNode :: ClassEnv -> TIPatternNode -> EvalM TIPatternNode
    expandTIPatternNode classEnv' node = case node of
      -- Loop pattern: expand the loop range expressions
      TILoopPat var loopRange pat1 pat2 -> do
        loopRange' <- expandTILoopRange classEnv' loopRange
        pat1' <- expandTIPattern classEnv' pat1
        pat2' <- expandTIPattern classEnv' pat2
        return $ TILoopPat var loopRange' pat1' pat2'

      -- Recursive pattern constructors
      TIAndPat pat1 pat2 -> do
        pat1' <- expandTIPattern classEnv' pat1
        pat2' <- expandTIPattern classEnv' pat2
        return $ TIAndPat pat1' pat2'

      TIOrPat pat1 pat2 -> do
        pat1' <- expandTIPattern classEnv' pat1
        pat2' <- expandTIPattern classEnv' pat2
        return $ TIOrPat pat1' pat2'

      TIForallPat pat1 pat2 -> do
        pat1' <- expandTIPattern classEnv' pat1
        pat2' <- expandTIPattern classEnv' pat2
        return $ TIForallPat pat1' pat2'

      TINotPat pat -> do
        pat' <- expandTIPattern classEnv' pat
        return $ TINotPat pat'

      TITuplePat pats -> do
        pats' <- mapM (expandTIPattern classEnv') pats
        return $ TITuplePat pats'

      TIInductivePat name pats -> do
        pats' <- mapM (expandTIPattern classEnv') pats
        return $ TIInductivePat name pats'

      TIIndexedPat pat exprs -> do
        pat' <- expandTIPattern classEnv' pat
        exprs' <- mapM (expandTIExprWithConstraints classEnv') exprs
        return $ TIIndexedPat pat' exprs'

      TILetPat bindings pat -> do
        pat' <- expandTIPattern classEnv' pat
        return $ TILetPat bindings pat'  -- TODO: Expand binding expressions
      
      TIPApplyPat funcExpr argPats -> do
        funcExpr' <- expandTIExprWithConstraints classEnv' funcExpr
        argPats' <- mapM (expandTIPattern classEnv') argPats
        return $ TIPApplyPat funcExpr' argPats'

      TIDApplyPat pat pats -> do
        pat' <- expandTIPattern classEnv' pat
        pats' <- mapM (expandTIPattern classEnv') pats
        return $ TIDApplyPat pat' pats'

      TISeqConsPat pat1 pat2 -> do
        pat1' <- expandTIPattern classEnv' pat1
        pat2' <- expandTIPattern classEnv' pat2
        return $ TISeqConsPat pat1' pat2'

      TISeqNilPat -> return TISeqNilPat

      TIVarPat name -> return $ TIVarPat name

      TIInductiveOrPApplyPat name pats -> do
        pats' <- mapM (expandTIPattern classEnv') pats
        return $ TIInductiveOrPApplyPat name pats'

      -- Leaf patterns: no expansion needed
      TIWildCard -> return TIWildCard
      TIPatVar name -> return $ TIPatVar name
      TIValuePat expr -> do
        expr' <- expandTIExprWithConstraints classEnv' expr
        return $ TIValuePat expr'
      TIPredPat pred -> do
        pred' <- expandTIExprWithConstraints classEnv' pred
        return $ TIPredPat pred'
      TIContPat -> return TIContPat
      TILaterPatVar -> return TILaterPatVar

    -- Expand loop range expressions (no parent constraints)
    expandTILoopRange :: ClassEnv -> TILoopRange -> EvalM TILoopRange
    expandTILoopRange classEnv' (TILoopRange start end rangePat) = do
      start' <- expandTIExprWithConstraints classEnv' start
      end' <- expandTIExprWithConstraints classEnv' end
      rangePat' <- expandTIPattern classEnv' rangePat
      return $ TILoopRange start' end' rangePat'

    -- Try to resolve a method call using type class constraints
    -- Dictionary passing: convert method calls to dictionary access
    tryResolveMethodCall :: ClassEnv -> [Constraint] -> String -> [TIExpr] -> EvalM (Maybe TIExprNode)
    tryResolveMethodCall classEnv' cs methodName expandedArgs = do
      -- Find a constraint that provides this method
      case findConstraintForMethod classEnv' methodName cs of
        Nothing -> return Nothing
        Just (Constraint className tyArg) -> do
          -- Look up the class to check if methodName is a method
          case lookupClass className classEnv' of
            Just classInfo -> do
              if methodName `elem` map fst (classMethods classInfo)
                then do
                  let methodKey = sanitizeMethodName methodName
                  -- Check if this is a type variable constraint
                  case tyArg of
                    TVar (TyVar _v) -> do
                      -- Type variable: use dictionary parameter
                      -- e.g., for {Eq a}, use dict_Eq (without type parameter)
                      typeEnv <- getTypeEnv
                      let dictParamName = "dict_" ++ className
                      -- Look up dictionary type from type environment
                      dictHashType <- case lookupEnv (stringToVar dictParamName) typeEnv of
                        Just (Forall _ _ dictType) -> return dictType
                        Nothing -> return $ THash TString TAny  -- Fallback
                      -- Get method type from ClassEnv instead of dictHashType
                      let methodType = getMethodTypeFromClass classEnv' className methodKey tyArg
                          -- No constraints: dictionary access resolves the constraint
                          methodScheme = Forall [] [] methodType
                          dictExpr = TIExpr (Forall [] [] dictHashType) (TIVarExpr dictParamName)
                          indexExpr = TIExpr (Forall [] [] TString) 
                                            (TIConstantExpr (StringExpr (pack methodKey)))
                          dictAccess = TIExpr methodScheme $
                                       TIIndexedExpr False dictExpr [Sub indexExpr]
                      -- Apply arguments: dictAccess arg1 arg2 ...
                      return $ Just $ TIApplyExpr dictAccess expandedArgs
                    _ -> do
                      -- Concrete type: try to find matching instance
                      let instances = lookupInstances className classEnv'
                      -- Use actual argument type if needed
                      let argTypes = map tiExprType expandedArgs
                          actualType = case (tyArg, argTypes) of
                            (TVar _, (t:_)) -> t  -- Use first argument's type
                            _ -> tyArg
                      -- Check if actualType is still a type variable
                      case actualType of
                        TVar (TyVar _v') -> do
                          -- Still a type variable: use dictionary parameter
                          typeEnv <- getTypeEnv
                          let dictParamName = "dict_" ++ className
                          -- Look up dictionary type from type environment
                          dictHashType <- case lookupEnv (stringToVar dictParamName) typeEnv of
                            Just (Forall _ _ dictType) -> return dictType
                            Nothing -> return $ THash TString TAny  -- Fallback
                          -- Get method type from ClassEnv instead of dictHashType
                          let methodType = getMethodTypeFromClass classEnv' className methodKey actualType
                              -- No constraints: dictionary access resolves the constraint
                              methodScheme = Forall [] [] methodType
                              dictExpr = TIExpr (Forall [] [] dictHashType) (TIVarExpr dictParamName)
                              indexExpr = TIExpr (Forall [] [] TString) 
                                                (TIConstantExpr (StringExpr (pack methodKey)))
                              dictAccess = TIExpr methodScheme $
                                           TIIndexedExpr False dictExpr [Sub indexExpr]
                          -- Apply arguments: dictAccess arg1 arg2 ...
                          return $ Just $ TIApplyExpr dictAccess expandedArgs
                        _ -> case findMatchingInstanceForType actualType instances of
                          Just inst -> do
                            -- Found an instance: generate dictionary access
                            -- e.g., numInteger_"plus" for Num Integer instance
                            typeEnv <- getTypeEnv
                            let instTypeName = typeConstructorName (instType inst)
                                dictName = lowerFirst className ++ instTypeName

                            -- Look up dictionary type from type environment
                            dictHashType <- case lookupEnv (stringToVar dictName) typeEnv of
                              Just (Forall _ _ dictType) -> return dictType
                              Nothing -> return $ THash TString TAny  -- Fallback

                            -- Get method type from ClassEnv instead of dictHashType
                            let methodType = getMethodTypeFromClass classEnv' className methodKey actualType
                                -- No constraints: dictionary access resolves the constraint
                                methodScheme = Forall [] [] methodType
                            
                            -- Check if instance has nested constraints
                            -- If so, dictionary is a function that takes dict parameters
                            dictExprBase <- if null (instContext inst)
                                  then do
                                    -- No constraints: dictionary is a simple hash
                                    let dictExpr = TIExpr (Forall [] [] dictHashType) (TIVarExpr dictName)
                                    return dictExpr
                                  else do
                                    -- Has constraints: dictionary is a function
                                    -- Need to resolve constraint arguments and apply them
                                    -- e.g., eqCollection eqInteger
                                    let dictFuncType = case dictHashType of
                                          TFun _ resultType -> TFun dictHashType resultType
                                          _ -> TFun (THash TString TAny) dictHashType
                                        dictFuncExpr = TIExpr (Forall [] [] dictFuncType) (TIVarExpr dictName)

                                    -- Substitute type variables in constraints with actual types
                                    -- e.g., for instance {Eq a} Eq [a] matched with [Integer]
                                    -- instType inst = [a], actualType = [Integer]
                                    -- constraint {Eq a} should become {Eq Integer}
                                    -- Substitute type variables in constraints
                                    -- e.g., instance {Eq a} Eq [a] matched with [[Integer]]
                                    -- instType = [a], actualType = [[Integer]]
                                    -- Extract a -> [Integer], apply to {Eq a} -> {Eq [Integer]}
                                    let substitutedConstraints = substituteInstanceConstraints (instType inst) actualType (instContext inst)
                                    -- Resolve each substituted constraint (depth is managed internally)
                                    dictArgs <- mapM (resolveDictionaryArg classEnv') substitutedConstraints
                                    -- Apply dictionary function to constraint dictionaries
                                    return $ TIExpr (Forall [] [] dictHashType) (TIApplyExpr dictFuncExpr dictArgs)

                                -- Now index into the dictionary (which is now a hash)
                            let indexExpr = TIExpr (Forall [] [] TString) 
                                                  (TIConstantExpr (StringExpr (pack methodKey)))
                                dictAccess = TIExpr methodScheme $
                                             TIIndexedExpr False dictExprBase [Sub indexExpr]
                            -- Apply arguments: dictAccess arg1 arg2 ...
                            return $ Just $ TIApplyExpr dictAccess expandedArgs
                          Nothing -> return Nothing
                else return Nothing
            Nothing -> return Nothing
    
    -- Substitute type variables in instance constraints based on actual type
    -- e.g., for instance {Eq a} Eq [a] matched with [[Integer]]
    -- instType = [a], actualType = [[Integer]]
    -- Extract: a -> [Integer], then apply to constraints {Eq a} -> {Eq [Integer]}
    substituteInstanceConstraints :: Type -> Type -> [Constraint] -> [Constraint]
    substituteInstanceConstraints instType actualType constraints =
      let substs = extractTypeSubstitutions instType actualType
      in map (applySubstsToConstraint substs) constraints

    -- Resolve a constraint to a dictionary argument (with depth limit to prevent infinite recursion)
    resolveDictionaryArg :: ClassEnv -> Constraint -> EvalM TIExpr
    resolveDictionaryArg classEnv constraint = resolveDictionaryArgWithDepth classEnv 50 constraint
    
    resolveDictionaryArgWithDepth :: ClassEnv -> Int -> Constraint -> EvalM TIExpr
    resolveDictionaryArgWithDepth _ 0 (Constraint className _) = do
      -- Depth limit reached, return error placeholder
      return $ TIExpr (Forall [] [] (TVar (TyVar "error"))) (TIVarExpr ("dict_" ++ className ++ "_TOO_DEEP"))
    
    resolveDictionaryArgWithDepth classEnv depth (Constraint className tyArg) = do
      case tyArg of
        TVar (TyVar _v) -> do
          -- Type variable: use dictionary parameter name (without type parameter)
          -- e.g., for {Eq a}, return dict_Eq
          let dictParamName = "dict_" ++ className
              dictType = TVar (TyVar "dict")
          return $ TIExpr (Forall [] [] dictType) (TIVarExpr dictParamName)
        _ -> do
          -- Concrete type: try to find matching instance
          let instances = lookupInstances className classEnv
          case findMatchingInstanceForType tyArg instances of
            Just inst -> do
              -- Found instance: generate dictionary name (e.g., "numInteger", "eqCollection")
              let instTypeName = typeConstructorName (instType inst)
                  dictName = lowerFirst className ++ instTypeName
                  dictType = TVar (TyVar "dict")
                  dictExpr = TIExpr (Forall [] [] dictType) (TIVarExpr dictName)
              
              -- Check if this instance has nested constraints
              -- e.g., instance {Eq a} Eq [a] has constraint {Eq a}
              if null (instContext inst)
                then do
                  -- No constraints: return simple dictionary reference
                  return dictExpr
                else do
                  -- Has constraints: need to resolve them and apply to dictionary
                  -- e.g., for Eq [Integer], resolve {Eq Integer} -> eqInteger
                  -- then return: eqCollection eqInteger

                  -- Substitute type variables in constraints with actual types
                  -- e.g., for instance {Eq a} Eq [a] matched with [[Integer]]
                  -- instType inst = [a], tyArg = [[Integer]]
                  -- Extract: a -> [Integer]
                  -- Apply to constraints: {Eq a} -> {Eq [Integer]}
                  let substs = extractTypeSubstitutions (instType inst) tyArg
                      substitutedConstraints = map (applySubstsToConstraint substs) (instContext inst)

                  -- Recursively resolve each constraint with reduced depth
                  dictArgs <- mapM (resolveDictionaryArgWithDepth classEnv (depth - 1)) substitutedConstraints

                  -- Apply dictionary function to resolved dictionaries
                  -- e.g., eqCollection eqInteger (when resolving Eq [Integer])
                  --       eqCollection (eqCollection eqInteger) (when resolving Eq [[Integer]])
                  return $ TIExpr (Forall [] [] dictType) (TIApplyExpr dictExpr dictArgs)
            Nothing -> do
              -- No instance found - this is an error, but return a dummy for now
              return $ TIExpr (Forall [] [] (TVar (TyVar "error"))) (TIVarExpr "undefined")

-- | Generate dictionary parameter name from constraint
-- Used for both dictionary parameter generation and dictionary argument passing
-- Type parameters are not included in the dictionary parameter name
constraintToDictParam :: Constraint -> String
constraintToDictParam (Constraint className _constraintType) =
  "dict_" ++ className

-- | Get method type from ClassEnv
-- This retrieves the method type from the class definition and substitutes type variables
-- Note: methodKey is the sanitized name (e.g., "plus"), but classMethods uses original names (e.g., "+")
-- We need to try both the sanitized and original names
getMethodTypeFromClass :: ClassEnv -> String -> String -> Type -> Type
getMethodTypeFromClass classEnv className methodKey constraintType =
  case lookupClass className classEnv of
    Just classInfo ->
      -- Try to find the method by sanitized name first, then try unsanitizing
      case lookup methodKey (classMethods classInfo) `mplus` lookupUnsanitized methodKey (classMethods classInfo) of
        Just classMethodType ->
          -- Substitute class type parameter with actual constraint type
          -- e.g., class Num a has plus : a -> a -> a
          --       constraint Num t0 → plus : t0 -> t0 -> t0
          applySubstsToType [(classParam classInfo, constraintType)] classMethodType
        Nothing -> TAny  -- Method not found in class
    Nothing -> TAny  -- Class not found
  where
    -- Lookup by unsanitizing the method key (reverse of sanitizeMethodName)
    -- e.g., "plus" -> "+", "times" -> "*"
    lookupUnsanitized :: String -> [(String, a)] -> Maybe a
    lookupUnsanitized key methods =
      case unsanitizeMethodName key of
        Just originalName -> lookup originalName methods
        Nothing -> Nothing

    -- Reverse of sanitizeMethodName
    unsanitizeMethodName :: String -> Maybe String
    unsanitizeMethodName "eq" = Just "=="
    unsanitizeMethodName "neq" = Just "/="
    unsanitizeMethodName "lt" = Just "<"
    unsanitizeMethodName "le" = Just "<="
    unsanitizeMethodName "gt" = Just ">"
    unsanitizeMethodName "ge" = Just ">="
    unsanitizeMethodName "plus" = Just "+"
    unsanitizeMethodName "minus" = Just "-"
    unsanitizeMethodName "times" = Just "*"
    unsanitizeMethodName "div" = Just "/"
    unsanitizeMethodName _ = Nothing

-- | Add dictionary parameters to a function based on its type scheme constraints
-- This transforms constrained functions into dictionary-passing style
addDictionaryParametersT :: TypeScheme -> TIExpr -> EvalM TIExpr
addDictionaryParametersT (Forall _vars constraints _ty) tiExpr
  | null constraints = return tiExpr  -- No constraints, no change
  | otherwise = do
      classEnv <- getClassEnv
      -- Note: No need to resolve Tensor constraints here because TensorMapInsertion
      -- runs before TypeClassExpand, so tensor operations are already handled.
      -- The execution order is: insertTensorMaps -> expandTypeClassMethodsT
      addDictParamsToTIExpr classEnv constraints tiExpr
  where
    -- Add dictionary parameters to a TIExpr
    addDictParamsToTIExpr :: ClassEnv -> [Constraint] -> TIExpr -> EvalM TIExpr
    addDictParamsToTIExpr env cs expr = case tiExprNode expr of
      -- Lambda: add dictionary parameters before regular parameters
      TILambdaExpr mVar params body -> do
        let dictParams = map constraintToDictParam cs
            dictVars = map stringToVar dictParams
        -- Replace method calls in body with dictionary access
        -- BUT: if body is a hash (dictionary), don't process it
        body' <- case tiExprNode body of
                   TIHashExpr _ -> return body  -- Dictionary body, don't process
                   _ -> replaceMethodCallsWithDictAccessT env cs body
        let newNode = TILambdaExpr mVar (dictVars ++ params) body'
        return $ TIExpr (tiScheme expr) newNode
      
      -- Hash (dictionary definition): wrap in lambda AND apply dict params to methods
      -- Dictionary values are method references that need dictionary parameters
      TIHashExpr pairs -> do
        let dictParams = map constraintToDictParam cs
            dictVars = map stringToVar dictParams
            wrapperType = tiExprType expr
        
        -- For each value in the hash (which is a method reference),
        -- if it has constraints, apply dictionary parameters to it
        pairs' <- mapM (\(k, v) -> do
          -- Check if the value (method) has constraints
          typeEnv <- getTypeEnv
          let vNode = tiExprNode v
          case vNode of
            TIVarExpr methodName -> do
              case lookupEnv (stringToVar methodName) typeEnv of
                Just (Forall _ vConstraints _) | not (null vConstraints) -> do
                  -- Method has constraints, apply dictionary parameters
                  let dictArgExprs = map (\p -> TIExpr (Forall [] [] (TVar (TyVar "dict"))) (TIVarExpr p)) dictParams
                      vApplied = TIExpr (tiScheme v) (TIApplyExpr v dictArgExprs)
                  return (k, vApplied)
                _ -> return (k, v)  -- No constraints, keep as-is
            _ -> return (k, v)  -- Not a variable, keep as-is
          ) pairs
        
        let hashExpr' = TIExpr (tiScheme expr) (TIHashExpr pairs')
            newNode = TILambdaExpr Nothing dictVars hashExpr'
            newScheme = Forall [] [] wrapperType
        return $ TIExpr newScheme newNode
      
      -- Not a lambda: wrap in a lambda with dictionary parameters
      _ -> do
        let dictParams = map constraintToDictParam cs
            dictVars = map stringToVar dictParams
        -- Special handling for TIVarExpr: if it's a constrained variable, apply dictionaries
        expr' <- case tiExprNode expr of
          TIVarExpr varName -> do
            -- Check if this variable has constraints that match our constraints
            typeEnv <- getTypeEnv
            case lookupEnv (stringToVar varName) typeEnv of
              Just (Forall _ varConstraints _) | not (null varConstraints) -> do
                -- Check which constraints from varConstraints match parent constraints cs
                let (Forall _ exprConstraints exprType) = tiScheme expr
                    matchingConstraints = filter (\(Constraint eName eType) ->
                          any (\(Constraint pName pType) ->
                            eName == pName && eType == pType) cs) exprConstraints
                if null matchingConstraints
                  then replaceMethodCallsWithDictAccessT env cs expr
                  else do
                    -- Apply matching dictionary parameters
                    let dictArgExprs = map (\p -> TIExpr (Forall [] [] (TVar (TyVar "dict"))) (TIVarExpr p))
                                           (map constraintToDictParam matchingConstraints)
                        varExpr = TIExpr (tiScheme expr) (TIVarExpr varName)
                    return $ TIExpr (tiScheme expr) (TIApplyExpr varExpr dictArgExprs)
              _ -> replaceMethodCallsWithDictAccessT env cs expr
          _ -> replaceMethodCallsWithDictAccessT env cs expr
        let wrapperType = tiExprType expr
            newNode = TILambdaExpr Nothing dictVars expr'
            newScheme = Forall [] [] wrapperType
        return $ TIExpr newScheme newNode
    
    -- Replace method calls with dictionary access in TIExpr
    replaceMethodCallsWithDictAccessT :: ClassEnv -> [Constraint] -> TIExpr -> EvalM TIExpr
    replaceMethodCallsWithDictAccessT env cs tiExpr = do
      let scheme@(Forall _ exprConstraints exprType) = tiScheme tiExpr
      newNode <- replaceMethodCallsInNode env cs exprConstraints exprType (tiExprNode tiExpr)
      return $ TIExpr scheme newNode
    
    -- Replace method calls in TIExprNode
    replaceMethodCallsInNode :: ClassEnv -> [Constraint] -> [Constraint] -> Type -> TIExprNode -> EvalM TIExprNode
    replaceMethodCallsInNode env cs exprConstraints exprType node = case node of
      -- Standalone method reference: eta-expand
      TIVarExpr methodName -> do
        case findConstraintForMethod env methodName cs of
          Just constraint -> do
            -- Get method type to determine arity
            typeEnv <- getTypeEnv
            case lookupEnv (stringToVar methodName) typeEnv of
              Just (Forall _ _ _ty) -> do
                -- Use the expression's actual type (exprType) instead of the method's declared type (ty)
                -- because eta-expansion should create parameters matching the expected usage context
                let arity = getMethodArity exprType
                    paramTypes = getParamTypes exprType
                    paramNames = ["etaVar" ++ show i | i <- [1..arity]]
                    paramVars = map stringToVar paramNames
                    paramExprs = zipWith (\n t -> TIExpr (Forall [] [] t) (TIVarExpr n)) paramNames paramTypes
                    -- Create dictionary access
                    dictParam = constraintToDictParam constraint
                    Constraint className tyArg = constraint
                -- Look up dictionary type from type environment
                dictHashType <- case lookupEnv (stringToVar dictParam) typeEnv of
                  Just (Forall _ _ dictType) -> return dictType
                  Nothing -> return $ THash TString TAny  -- Fallback
                -- Get method type from ClassEnv instead of dictHashType
                let methodType = getMethodTypeFromClass env className (sanitizeMethodName methodName) tyArg
                    methodConstraint = Constraint className tyArg
                    methodScheme = Forall (Set.toList $ freeTyVars tyArg) [methodConstraint] methodType
                    indexExpr = TIExpr (Forall [] [] TString) 
                                      (TIConstantExpr (StringExpr (pack (sanitizeMethodName methodName))))
                    dictAccess = TIExpr methodScheme $
                                 TIIndexedExpr False
                                   (TIExpr (Forall [] [] dictHashType) (TIVarExpr dictParam))
                                   [Sub indexExpr]
                    -- Create: dictAccess etaVar1 etaVar2 ... etaVarN
                    body = TIExpr methodScheme (TIApplyExpr dictAccess paramExprs)
                return $ TILambdaExpr Nothing paramVars body
              Nothing -> return $ TIVarExpr methodName
          Nothing -> do
            -- Not a method - just return the variable as-is
            -- Dictionary application for constrained variables is handled by expandTypeClassMethodsT
            return $ TIVarExpr methodName
      
      -- Method call: replace with dictionary access
      TIApplyExpr func args -> do
        case tiExprNode func of
          TIVarExpr methodName -> do
            case findConstraintForMethod env methodName cs of
              Just constraint -> do
                -- Replace with dictionary access
                typeEnv <- getTypeEnv
                let dictParam = constraintToDictParam constraint
                    Constraint className tyArg = constraint
                -- Look up dictionary type from type environment
                dictHashType <- case lookupEnv (stringToVar dictParam) typeEnv of
                  Just (Forall _ _ dictType) -> return dictType
                  Nothing -> return $ THash TString TAny  -- Fallback
                -- Get method type from ClassEnv instead of dictHashType
                let methodType = getMethodTypeFromClass env className (sanitizeMethodName methodName) tyArg
                    methodConstraint = Constraint className tyArg
                    methodScheme = Forall (Set.toList $ freeTyVars tyArg) [methodConstraint] methodType
                    indexExpr = TIExpr (Forall [] [] TString) 
                                      (TIConstantExpr (StringExpr (pack (sanitizeMethodName methodName))))
                    dictAccessNode = TIIndexedExpr False
                                     (TIExpr (Forall [] [] dictHashType) (TIVarExpr dictParam))
                                     [Sub indexExpr]
                    dictAccess = TIExpr methodScheme dictAccessNode
                -- Recursively process arguments
                args' <- mapM (replaceMethodCallsWithDictAccessT env cs) args
                return $ TIApplyExpr dictAccess args'
              Nothing -> do
                -- Not a method, process recursively
                func' <- replaceMethodCallsWithDictAccessT env cs func
                args' <- mapM (replaceMethodCallsWithDictAccessT env cs) args
                return $ TIApplyExpr func' args'
          _ -> do
            -- Not a simple variable, process recursively
            func' <- replaceMethodCallsWithDictAccessT env cs func
            args' <- mapM (replaceMethodCallsWithDictAccessT env cs) args
            return $ TIApplyExpr func' args'
      
      -- Lambda: recursively process body
      TILambdaExpr mVar params body -> do
        body' <- replaceMethodCallsWithDictAccessT env cs body
        return $ TILambdaExpr mVar params body'
      
      -- If: recursively process
      TIIfExpr cond thenExpr elseExpr -> do
        cond' <- replaceMethodCallsWithDictAccessT env cs cond
        thenExpr' <- replaceMethodCallsWithDictAccessT env cs thenExpr
        elseExpr' <- replaceMethodCallsWithDictAccessT env cs elseExpr
        return $ TIIfExpr cond' thenExpr' elseExpr'
      
      -- Let: recursively process
      TILetExpr bindings body -> do
        bindings' <- mapM (\(pat, e) -> do
          e' <- replaceMethodCallsWithDictAccessT env cs e
          return (pat, e')) bindings
        body' <- replaceMethodCallsWithDictAccessT env cs body
        return $ TILetExpr bindings' body'
      
      -- LetRec: recursively process
      TILetRecExpr bindings body -> do
        bindings' <- mapM (\(pat, e) -> do
          e' <- replaceMethodCallsWithDictAccessT env cs e
          return (pat, e')) bindings
        body' <- replaceMethodCallsWithDictAccessT env cs body
        return $ TILetRecExpr bindings' body'
      
      -- Hash: do NOT process values inside dictionary hashes
      -- Dictionary values should remain as simple references
      -- e.g., {| ("eq", eqCollectionEq), ... |} not {| ("eq", eqCollectionEq dict_Eq), ... |}
      -- We return the node as-is without recursively processing the pairs
      TIHashExpr pairs -> do
        -- Process only keys, not values (values should remain as method references)
        pairs' <- mapM (\(k, v) -> do
          k' <- replaceMethodCallsWithDictAccessT env cs k
          -- Do NOT process v - keep it as a simple reference
          return (k', v)) pairs
        return $ TIHashExpr pairs'
      
      -- Matcher: recursively process expressions inside matcher definitions
      TIMatcherExpr patDefs -> do
        patDefs' <- mapM (\(pat, matcherExpr, bindings) -> do
          -- Process the next-matcher expression
          matcherExpr' <- replaceMethodCallsWithDictAccessT env cs matcherExpr
          -- Process expressions in primitive-data-match clauses
          bindings' <- mapM (\(dp, expr) -> do
            expr' <- replaceMethodCallsWithDictAccessT env cs expr
            return (dp, expr')) bindings
          return (pat, matcherExpr', bindings')) patDefs
        return $ TIMatcherExpr patDefs'
      
      -- Other expressions: return as-is for now
      _ -> return node

-- | Apply dictionaries to expressions with concrete type constraints
-- This is used for top-level definitions like: def integer : Matcher Integer := eq
-- where the right-hand side (eq) has concrete type constraints {Eq Integer}
applyConcreteConstraintDictionaries :: TIExpr -> EvalM TIExpr
applyConcreteConstraintDictionaries expr = do
  classEnv <- getClassEnv
  let scheme@(Forall vars constraints _) = tiScheme expr

  -- First, recursively process sub-expressions
  expr' <- case tiExprNode expr of
    TIApplyExpr func args -> do
      func' <- applyConcreteConstraintDictionaries func
      args' <- mapM applyConcreteConstraintDictionaries args
      return $ TIExpr scheme (TIApplyExpr func' args')
    _ -> return expr

  -- Then check if this expression has concrete constraints
  let isConcreteConstraint (Constraint _ (TVar _)) = False
      isConcreteConstraint _ = True
      hasOnlyConcreteConstraints = not (null constraints) && all isConcreteConstraint constraints

  if hasOnlyConcreteConstraints
    then do
      -- Apply dictionaries for concrete constraints
      dictArgs <- mapM (resolveDictionaryForConstraint classEnv) constraints
      -- Create application: expr dict1 dict2 ...
      let resultType = tiExprType expr'
          -- Update scheme to remove constraints since they are now applied
          -- Keep type variables (vars) as they may be needed for polymorphism
          newScheme = Forall vars [] resultType
      return $ TIExpr newScheme (TIApplyExpr expr' dictArgs)
    else
      -- No concrete constraints, return as-is
      return expr'
  where
    -- Resolve dictionary for a concrete constraint
    resolveDictionaryForConstraint :: ClassEnv -> Constraint -> EvalM TIExpr
    resolveDictionaryForConstraint classEnv (Constraint className tyArg) = do
      -- Normalize TInt to TMathExpr for instance matching
      -- Integer and MathExpr are the same type in Egison
      let normalizedType = case tyArg of
                             TInt -> TMathExpr
                             _ -> tyArg
      let instances = lookupInstances className classEnv
      case findMatchingInstanceForType normalizedType instances of
        Just inst -> do
          -- Generate dictionary name (e.g., "eqInteger", "numInteger")
          let instTypeName = typeConstructorName (instType inst)
              dictName = lowerFirst className ++ instTypeName
              dictType = TVar (TyVar "dict")
              dictExpr = TIExpr (Forall [] [] dictType) (TIVarExpr dictName)
          
          -- Check if instance has nested constraints
          if null (instContext inst)
            then do
              -- No constraints: return simple dictionary reference
              return dictExpr
            else do
              -- Has constraints: need to resolve them recursively
              nestedDictArgs <- mapM (resolveDictionaryForConstraint classEnv) (instContext inst)
              return $ TIExpr (Forall [] [] dictType) (TIApplyExpr dictExpr nestedDictArgs)
        Nothing -> do
          -- No instance found - return dummy dictionary
          let dictName = "dict_" ++ className ++ "_NOT_FOUND"
              dictType = TVar (TyVar "dict")
          return $ TIExpr (Forall [] [] dictType) (TIVarExpr dictName)

-- | Expand type class method calls in patterns
-- This is a public wrapper for expandTIPattern used by TypedDesugar
expandTypeClassMethodsInPattern :: TIPattern -> EvalM TIPattern
expandTypeClassMethodsInPattern tipat = do
  classEnv <- getClassEnv
  expandPatternWithClassEnv classEnv tipat
  where
    expandPatternWithClassEnv :: ClassEnv -> TIPattern -> EvalM TIPattern
    expandPatternWithClassEnv classEnv' (TIPattern scheme node) = do
      node' <- expandPatternNode classEnv' node
      return $ TIPattern scheme node'
    
    expandPatternNode :: ClassEnv -> TIPatternNode -> EvalM TIPatternNode
    expandPatternNode classEnv' node = case node of
      TILoopPat var loopRange pat1 pat2 -> do
        loopRange' <- expandLoopRange classEnv' loopRange
        pat1' <- expandPatternWithClassEnv classEnv' pat1
        pat2' <- expandPatternWithClassEnv classEnv' pat2
        return $ TILoopPat var loopRange' pat1' pat2'
      
      TIAndPat pat1 pat2 -> do
        pat1' <- expandPatternWithClassEnv classEnv' pat1
        pat2' <- expandPatternWithClassEnv classEnv' pat2
        return $ TIAndPat pat1' pat2'
      
      TIOrPat pat1 pat2 -> do
        pat1' <- expandPatternWithClassEnv classEnv' pat1
        pat2' <- expandPatternWithClassEnv classEnv' pat2
        return $ TIOrPat pat1' pat2'
      
      TIForallPat pat1 pat2 -> do
        pat1' <- expandPatternWithClassEnv classEnv' pat1
        pat2' <- expandPatternWithClassEnv classEnv' pat2
        return $ TIForallPat pat1' pat2'
      
      TINotPat pat -> do
        pat' <- expandPatternWithClassEnv classEnv' pat
        return $ TINotPat pat'
      
      TITuplePat pats -> do
        pats' <- mapM (expandPatternWithClassEnv classEnv') pats
        return $ TITuplePat pats'
      
      TIInductivePat name pats -> do
        pats' <- mapM (expandPatternWithClassEnv classEnv') pats
        return $ TIInductivePat name pats'
      
      TIIndexedPat pat exprs -> do
        pat' <- expandPatternWithClassEnv classEnv' pat
        exprs' <- mapM expandTypeClassMethodsT exprs
        return $ TIIndexedPat pat' exprs'
      
      TILetPat bindings pat -> do
        pat' <- expandPatternWithClassEnv classEnv' pat
        bindings' <- mapM (\(pd, e) -> do
          e' <- expandTypeClassMethodsT e
          return (pd, e')) bindings
        return $ TILetPat bindings' pat'
      
      TIPApplyPat funcExpr argPats -> do
        funcExpr' <- expandTypeClassMethodsT funcExpr
        argPats' <- mapM (expandPatternWithClassEnv classEnv') argPats
        return $ TIPApplyPat funcExpr' argPats'
      
      TIDApplyPat pat pats -> do
        pat' <- expandPatternWithClassEnv classEnv' pat
        pats' <- mapM (expandPatternWithClassEnv classEnv') pats
        return $ TIDApplyPat pat' pats'
      
      TISeqConsPat pat1 pat2 -> do
        pat1' <- expandPatternWithClassEnv classEnv' pat1
        pat2' <- expandPatternWithClassEnv classEnv' pat2
        return $ TISeqConsPat pat1' pat2'
      
      TIInductiveOrPApplyPat name pats -> do
        pats' <- mapM (expandPatternWithClassEnv classEnv') pats
        return $ TIInductiveOrPApplyPat name pats'
      
      TIValuePat expr -> do
        expr' <- expandTypeClassMethodsT expr
        expr'' <- applyConcreteConstraintDictionaries expr'
        return $ TIValuePat expr''
      
      TIPredPat pred -> do
        pred' <- expandTypeClassMethodsT pred
        pred'' <- applyConcreteConstraintDictionaries pred'
        return $ TIPredPat pred''
      
      -- Leaf patterns
      TISeqNilPat -> return TISeqNilPat
      TIVarPat name -> return $ TIVarPat name
      TIWildCard -> return TIWildCard
      TIPatVar name -> return $ TIPatVar name
      TIContPat -> return TIContPat
      TILaterPatVar -> return TILaterPatVar
    
    expandLoopRange :: ClassEnv -> TILoopRange -> EvalM TILoopRange
    expandLoopRange classEnv' (TILoopRange start end rangePat) = do
      start' <- expandTypeClassMethodsT start
      end' <- expandTypeClassMethodsT end
      rangePat' <- expandPatternWithClassEnv classEnv' rangePat
      return $ TILoopRange start' end' rangePat'

-- | Apply dictionaries to expressions with concrete constraints in patterns
-- This is used to apply dictionaries to value patterns like #(n + 1)
applyConcreteConstraintDictionariesInPattern :: TIPattern -> EvalM TIPattern
applyConcreteConstraintDictionariesInPattern (TIPattern scheme node) = do
  node' <- applyDictInPatternNode node
  return $ TIPattern scheme node'
  where
    applyDictInPatternNode :: TIPatternNode -> EvalM TIPatternNode
    applyDictInPatternNode pnode = case pnode of
      TIValuePat expr -> do
        expr' <- applyConcreteConstraintDictionaries expr
        return $ TIValuePat expr'
      
      TIPredPat expr -> do
        expr' <- applyConcreteConstraintDictionaries expr
        return $ TIPredPat expr'
      
      TIIndexedPat pat exprs -> do
        pat' <- applyConcreteConstraintDictionariesInPattern pat
        exprs' <- mapM applyConcreteConstraintDictionaries exprs
        return $ TIIndexedPat pat' exprs'
      
      TILetPat bindings pat -> do
        pat' <- applyConcreteConstraintDictionariesInPattern pat
        bindings' <- mapM (\(pd, e) -> do
          e' <- applyConcreteConstraintDictionaries e
          return (pd, e')) bindings
        return $ TILetPat bindings' pat'
      
      TILoopPat var loopRange pat1 pat2 -> do
        loopRange' <- applyDictInLoopRange loopRange
        pat1' <- applyConcreteConstraintDictionariesInPattern pat1
        pat2' <- applyConcreteConstraintDictionariesInPattern pat2
        return $ TILoopPat var loopRange' pat1' pat2'
      
      TIAndPat pat1 pat2 -> do
        pat1' <- applyConcreteConstraintDictionariesInPattern pat1
        pat2' <- applyConcreteConstraintDictionariesInPattern pat2
        return $ TIAndPat pat1' pat2'
      
      TIOrPat pat1 pat2 -> do
        pat1' <- applyConcreteConstraintDictionariesInPattern pat1
        pat2' <- applyConcreteConstraintDictionariesInPattern pat2
        return $ TIOrPat pat1' pat2'
      
      TIForallPat pat1 pat2 -> do
        pat1' <- applyConcreteConstraintDictionariesInPattern pat1
        pat2' <- applyConcreteConstraintDictionariesInPattern pat2
        return $ TIForallPat pat1' pat2'
      
      TINotPat pat -> do
        pat' <- applyConcreteConstraintDictionariesInPattern pat
        return $ TINotPat pat'
      
      TITuplePat pats -> do
        pats' <- mapM applyConcreteConstraintDictionariesInPattern pats
        return $ TITuplePat pats'
      
      TIInductivePat name pats -> do
        pats' <- mapM applyConcreteConstraintDictionariesInPattern pats
        return $ TIInductivePat name pats'
      
      TIPApplyPat funcExpr argPats -> do
        funcExpr' <- applyConcreteConstraintDictionaries funcExpr
        argPats' <- mapM applyConcreteConstraintDictionariesInPattern argPats
        return $ TIPApplyPat funcExpr' argPats'
      
      TIDApplyPat pat pats -> do
        pat' <- applyConcreteConstraintDictionariesInPattern pat
        pats' <- mapM applyConcreteConstraintDictionariesInPattern pats
        return $ TIDApplyPat pat' pats'
      
      TISeqConsPat pat1 pat2 -> do
        pat1' <- applyConcreteConstraintDictionariesInPattern pat1
        pat2' <- applyConcreteConstraintDictionariesInPattern pat2
        return $ TISeqConsPat pat1' pat2'
      
      TIInductiveOrPApplyPat name pats -> do
        pats' <- mapM applyConcreteConstraintDictionariesInPattern pats
        return $ TIInductiveOrPApplyPat name pats'
      
      -- Leaf patterns
      TISeqNilPat -> return TISeqNilPat
      TIVarPat name -> return $ TIVarPat name
      TIWildCard -> return TIWildCard
      TIPatVar name -> return $ TIPatVar name
      TIContPat -> return TIContPat
      TILaterPatVar -> return TILaterPatVar
    
    applyDictInLoopRange :: TILoopRange -> EvalM TILoopRange
    applyDictInLoopRange (TILoopRange start end rangePat) = do
      start' <- applyConcreteConstraintDictionaries start
      end' <- applyConcreteConstraintDictionaries end
      rangePat' <- applyConcreteConstraintDictionariesInPattern rangePat
      return $ TILoopRange start' end' rangePat'