{- |
Module      : Language.Egison.EnvBuilder
Licence     : MIT

This module implements Phase 2: Environment Building Phase.
It collects all declarations from TopExpr list before type inference and evaluation.

Environment Building Phase (Phase 2):
  1. Data constructor definitions collection (from InductiveDecl)
  2. Type class definitions collection (from ClassDeclExpr)
  3. Instance definitions collection (from InstanceDeclExpr)
  4. Type signature collection (from Define, DefineWithType)

This phase must be completed BEFORE type inference (Phase 5) begins,
ensuring all necessary information is available for type checking.
-}

module Language.Egison.EnvBuilder
  ( buildEnvironments
  , EnvBuildResult(..)
  ) where

import           Control.Monad              (foldM)
import           Control.Monad.Except       (throwError)
import           Control.Monad.State
import           Data.Char                  (toUpper, toLower)
import qualified Data.HashMap.Strict        as HashMap

import           Language.Egison.AST
import           Language.Egison.Data       (EvalM)
import           Language.Egison.EvalState  (ConstructorInfo(..), ConstructorEnv, PatternConstructorEnv)
import           Language.Egison.IExpr      (Var(..), Index(..), stringToVar)
import           Language.Egison.Desugar    (transVarIndex)
import           Language.Egison.Type.Env   (TypeEnv, ClassEnv, PatternTypeEnv, emptyEnv, emptyClassEnv, emptyPatternEnv,
                                             extendEnv, extendPatternEnv, addClass, addInstance, lookupClass)
import qualified Language.Egison.Type.Types as Types
import           Language.Egison.Type.Types (Type(..), TyVar(..), Constraint(..), TypeScheme(..), TensorShape(..),
                                             ClassInfo, InstanceInfo, freeTyVars, typeToName, sanitizeMethodName, typeExprToType,
                                             capitalizeFirst, lowerFirst)
import qualified Data.Set as Set

-- | Result of environment building phase
data EnvBuildResult = EnvBuildResult
  { EnvBuildResult -> TypeEnv
ebrTypeEnv        :: TypeEnv         -- ^ Type signatures for definitions
  , EnvBuildResult -> ClassEnv
ebrClassEnv       :: ClassEnv        -- ^ Type class and instance information
  , EnvBuildResult -> ConstructorEnv
ebrConstructorEnv :: ConstructorEnv  -- ^ Data constructor information
  , EnvBuildResult -> PatternConstructorEnv
ebrPatternConstructorEnv :: PatternConstructorEnv  -- ^ Pattern constructor information
  , EnvBuildResult -> PatternConstructorEnv
ebrPatternTypeEnv :: PatternTypeEnv  -- ^ Pattern function information
  } deriving (Int -> EnvBuildResult -> ShowS
[EnvBuildResult] -> ShowS
EnvBuildResult -> String
(Int -> EnvBuildResult -> ShowS)
-> (EnvBuildResult -> String)
-> ([EnvBuildResult] -> ShowS)
-> Show EnvBuildResult
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> EnvBuildResult -> ShowS
showsPrec :: Int -> EnvBuildResult -> ShowS
$cshow :: EnvBuildResult -> String
show :: EnvBuildResult -> String
$cshowList :: [EnvBuildResult] -> ShowS
showList :: [EnvBuildResult] -> ShowS
Show)

--------------------------------------------------------------------------------
-- Phase 2: Environment Building Phase
--------------------------------------------------------------------------------

-- | Build all environments from a list of top-level expressions.
-- This function implements Phase 2 of the processing flow.
-- It must be called AFTER expandLoads (Phase 1) and BEFORE type inference (Phase 5).
buildEnvironments :: [TopExpr] -> EvalM EnvBuildResult
buildEnvironments :: [TopExpr] -> EvalM EnvBuildResult
buildEnvironments [TopExpr]
exprs = do
  -- Start with empty environments
  let initialResult :: EnvBuildResult
initialResult = EnvBuildResult
        { ebrTypeEnv :: TypeEnv
ebrTypeEnv = TypeEnv
emptyEnv
        , ebrClassEnv :: ClassEnv
ebrClassEnv = ClassEnv
emptyClassEnv
        , ebrConstructorEnv :: ConstructorEnv
ebrConstructorEnv = ConstructorEnv
forall k v. HashMap k v
HashMap.empty
        , ebrPatternConstructorEnv :: PatternConstructorEnv
ebrPatternConstructorEnv = PatternConstructorEnv
emptyPatternEnv
        , ebrPatternTypeEnv :: PatternConstructorEnv
ebrPatternTypeEnv = PatternConstructorEnv
emptyPatternEnv
        }
  
  -- Process each top-level expression to collect declarations
  (EnvBuildResult -> TopExpr -> EvalM EnvBuildResult)
-> EnvBuildResult -> [TopExpr] -> EvalM EnvBuildResult
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM EnvBuildResult -> TopExpr -> EvalM EnvBuildResult
processTopExpr EnvBuildResult
initialResult [TopExpr]
exprs

-- | Process a single top-level expression to collect environment information
processTopExpr :: EnvBuildResult -> TopExpr -> EvalM EnvBuildResult
processTopExpr :: EnvBuildResult -> TopExpr -> EvalM EnvBuildResult
processTopExpr EnvBuildResult
result TopExpr
topExpr = case TopExpr
topExpr of
  
  -- 1. Data Constructor Definitions (from InductiveDecl)
  InductiveDecl String
typeName [String]
typeParams [InductiveConstructor]
constructors -> do
    let typeParamVars :: [Type]
typeParamVars = (String -> Type) -> [String] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map (TyVar -> Type
TVar (TyVar -> Type) -> (String -> TyVar) -> String -> Type
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> TyVar
TyVar) [String]
typeParams
        adtType :: Type
adtType = String -> [Type] -> Type
TInductive String
typeName [Type]
typeParamVars
        typeEnv :: TypeEnv
typeEnv = EnvBuildResult -> TypeEnv
ebrTypeEnv EnvBuildResult
result
        ctorEnv :: ConstructorEnv
ctorEnv = EnvBuildResult -> ConstructorEnv
ebrConstructorEnv EnvBuildResult
result
    
    -- Register each constructor
    (TypeEnv
typeEnv', ConstructorEnv
ctorEnv') <- ((TypeEnv, ConstructorEnv)
 -> InductiveConstructor
 -> StateT
      EvalState (ExceptT EgisonError RuntimeM) (TypeEnv, ConstructorEnv))
-> (TypeEnv, ConstructorEnv)
-> [InductiveConstructor]
-> StateT
     EvalState (ExceptT EgisonError RuntimeM) (TypeEnv, ConstructorEnv)
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM (String
-> [String]
-> Type
-> (TypeEnv, ConstructorEnv)
-> InductiveConstructor
-> StateT
     EvalState (ExceptT EgisonError RuntimeM) (TypeEnv, ConstructorEnv)
registerConstructor String
typeName [String]
typeParams Type
adtType) 
                                   (TypeEnv
typeEnv, ConstructorEnv
ctorEnv) 
                                   [InductiveConstructor]
constructors
    
    EnvBuildResult -> EvalM EnvBuildResult
forall a. a -> StateT EvalState (ExceptT EgisonError RuntimeM) a
forall (m :: * -> *) a. Monad m => a -> m a
return EnvBuildResult
result { ebrTypeEnv = typeEnv', ebrConstructorEnv = ctorEnv' }
  
  -- 2. Type Class Definitions (from ClassDeclExpr)
  ClassDeclExpr (ClassDecl String
className [String
typeParam] [ConstraintExpr]
superClasses [ClassMethod]
methods) -> do
    let classEnv :: ClassEnv
classEnv = EnvBuildResult -> ClassEnv
ebrClassEnv EnvBuildResult
result
        typeEnv :: TypeEnv
typeEnv = EnvBuildResult -> TypeEnv
ebrTypeEnv EnvBuildResult
result
        tyVar :: TyVar
tyVar = String -> TyVar
TyVar String
typeParam
        
        -- Extract superclass names from ConstraintExprs
        superNames :: [String]
superNames = (ConstraintExpr -> String) -> [ConstraintExpr] -> [String]
forall a b. (a -> b) -> [a] -> [b]
map ConstraintExpr -> String
extractConstraintName [ConstraintExpr]
superClasses
        
        -- Build method list with types
        methodsWithTypes :: [(String, Type)]
methodsWithTypes = (ClassMethod -> (String, Type))
-> [ClassMethod] -> [(String, Type)]
forall a b. (a -> b) -> [a] -> [b]
map ClassMethod -> (String, Type)
extractMethodWithType [ClassMethod]
methods
        
        -- Create ClassInfo
        -- Note: Use qualified name to avoid ambiguity with ClassDecl.classMethods
        classInfo :: ClassInfo
classInfo = Types.ClassInfo
          { classSupers :: [String]
Types.classSupers = [String]
superNames
          , classParam :: TyVar
Types.classParam = TyVar
tyVar
          , classMethods :: [(String, Type)]
Types.classMethods = [(String, Type)]
methodsWithTypes
          }
        
        -- Register class
        classEnv' :: ClassEnv
classEnv' = String -> ClassInfo -> ClassEnv -> ClassEnv
addClass String
className ClassInfo
classInfo ClassEnv
classEnv
        
        -- Register each class method to type environment
        typeEnv' :: TypeEnv
typeEnv' = (TypeEnv -> ClassMethod -> TypeEnv)
-> TypeEnv -> [ClassMethod] -> TypeEnv
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (TyVar -> String -> TypeEnv -> ClassMethod -> TypeEnv
registerClassMethod TyVar
tyVar String
className) TypeEnv
typeEnv [ClassMethod]
methods
    
    EnvBuildResult -> EvalM EnvBuildResult
forall a. a -> StateT EvalState (ExceptT EgisonError RuntimeM) a
forall (m :: * -> *) a. Monad m => a -> m a
return EnvBuildResult
result { ebrClassEnv = classEnv', ebrTypeEnv = typeEnv' }
  
  ClassDeclExpr ClassDecl
_ -> 
    -- Unsupported class declaration format (multiple type parameters, etc.)
    EnvBuildResult -> EvalM EnvBuildResult
forall a. a -> StateT EvalState (ExceptT EgisonError RuntimeM) a
forall (m :: * -> *) a. Monad m => a -> m a
return EnvBuildResult
result
  
  -- 3. Instance Definitions (from InstanceDeclExpr)
  InstanceDeclExpr (InstanceDecl [ConstraintExpr]
context String
className [TypeExpr]
instTypes [InstanceMethod]
methods) -> do
    let classEnv :: ClassEnv
classEnv = EnvBuildResult -> ClassEnv
ebrClassEnv EnvBuildResult
result
        typeEnv :: TypeEnv
typeEnv = EnvBuildResult -> TypeEnv
ebrTypeEnv EnvBuildResult
result
        
        -- Get the main instance type
        mainInstType :: Type
mainInstType = case [TypeExpr]
instTypes of
          []    -> Type
TAny
          (TypeExpr
t:[TypeExpr]
_) -> TypeExpr -> Type
typeExprToType TypeExpr
t
        
        -- Create InstanceInfo
        instInfo :: InstanceInfo
instInfo = Types.InstanceInfo
          { instContext :: [Constraint]
Types.instContext = (ConstraintExpr -> Constraint) -> [ConstraintExpr] -> [Constraint]
forall a b. (a -> b) -> [a] -> [b]
map ConstraintExpr -> Constraint
constraintToInternal [ConstraintExpr]
context
          , instClass :: String
Types.instClass = String
className
          , instType :: Type
Types.instType = Type
mainInstType
          , instMethods :: [(String, ())]
Types.instMethods = []  -- Methods are handled during desugaring/evaluation
          }
        
        -- Register instance
        classEnv' :: ClassEnv
classEnv' = String -> InstanceInfo -> ClassEnv -> ClassEnv
addInstance String
className InstanceInfo
instInfo ClassEnv
classEnv
        
        -- Register method type signatures for generated methods
        -- This prevents "Unbound variable" warnings during type inference
        -- Pass the instance context (constraints) to include in method types
        typeEnv' :: TypeEnv
typeEnv' = String
-> Type
-> [Constraint]
-> [InstanceMethod]
-> ClassEnv
-> TypeEnv
-> TypeEnv
registerInstanceMethods String
className Type
mainInstType ((ConstraintExpr -> Constraint) -> [ConstraintExpr] -> [Constraint]
forall a b. (a -> b) -> [a] -> [b]
map ConstraintExpr -> Constraint
constraintToInternal [ConstraintExpr]
context) [InstanceMethod]
methods ClassEnv
classEnv' TypeEnv
typeEnv
    
    EnvBuildResult -> EvalM EnvBuildResult
forall a. a -> StateT EvalState (ExceptT EgisonError RuntimeM) a
forall (m :: * -> *) a. Monad m => a -> m a
return EnvBuildResult
result { ebrClassEnv = classEnv', ebrTypeEnv = typeEnv' }
  
  -- 4. Type Signature Collection (from Define, DefineWithType)
  -- Note: We only collect explicit type signatures here.
  -- Inferred types will be added during type inference.
  DefineWithType TypedVarWithIndices
typedVar Expr
_expr -> do
    let name :: String
name = TypedVarWithIndices -> String
typedVarName TypedVarWithIndices
typedVar
        varIndices :: [VarIndex]
varIndices = TypedVarWithIndices -> [VarIndex]
typedVarIndices TypedVarWithIndices
typedVar
        -- Convert VarIndex to Index (Maybe Var) - like transVarIndex but with Nothing content
        indexTypes :: [Index (Maybe Var)]
indexTypes = (VarIndex -> [Index (Maybe Var)])
-> [VarIndex] -> [Index (Maybe Var)]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap VarIndex -> [Index (Maybe Var)]
transVarIndex [VarIndex]
varIndices
        -- Create Var with index structure (content is Just Var, so map to Nothing)
        var :: Var
var = String -> [Index (Maybe Var)] -> Var
Var String
name ((Index (Maybe Var) -> Index (Maybe Var))
-> [Index (Maybe Var)] -> [Index (Maybe Var)]
forall a b. (a -> b) -> [a] -> [b]
map ((Maybe Var -> Maybe Var) -> Index (Maybe Var) -> Index (Maybe Var)
forall a b. (a -> b) -> Index a -> Index b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Maybe Var -> Maybe Var -> Maybe Var
forall a b. a -> b -> a
const Maybe Var
forall a. Maybe a
Nothing)) [Index (Maybe Var)]
indexTypes)
        params :: [TypedParam]
params = TypedVarWithIndices -> [TypedParam]
typedVarParams TypedVarWithIndices
typedVar
        retType :: Type
retType = TypeExpr -> Type
typeExprToType (TypedVarWithIndices -> TypeExpr
typedVarRetType TypedVarWithIndices
typedVar)
        paramTypes :: [Type]
paramTypes = (TypedParam -> Type) -> [TypedParam] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map TypedParam -> Type
typedParamToType [TypedParam]
params
        
        -- Build function type
        funType :: Type
funType = (Type -> Type -> Type) -> Type -> [Type] -> Type
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr Type -> Type -> Type
TFun Type
retType [Type]
paramTypes
        
        -- Convert constraints from AST to internal representation
        constraints :: [Constraint]
constraints = (ConstraintExpr -> Constraint) -> [ConstraintExpr] -> [Constraint]
forall a b. (a -> b) -> [a] -> [b]
map ConstraintExpr -> Constraint
constraintToInternal (TypedVarWithIndices -> [ConstraintExpr]
typedVarConstraints TypedVarWithIndices
typedVar)
        
        -- Generalize free type variables in the type signature
        -- This handles type parameters like {a, b, c} in def compose {a, b, c} ...
        freeVars :: [TyVar]
freeVars = Set TyVar -> [TyVar]
forall a. Set a -> [a]
Set.toList (Type -> Set TyVar
freeTyVars Type
funType)
        typeScheme :: TypeScheme
typeScheme = [TyVar] -> [Constraint] -> Type -> TypeScheme
Types.Forall [TyVar]
freeVars [Constraint]
constraints Type
funType
        
        typeEnv :: TypeEnv
typeEnv = EnvBuildResult -> TypeEnv
ebrTypeEnv EnvBuildResult
result
        typeEnv' :: TypeEnv
typeEnv' = Var -> TypeScheme -> TypeEnv -> TypeEnv
extendEnv Var
var TypeScheme
typeScheme TypeEnv
typeEnv
    
    EnvBuildResult -> EvalM EnvBuildResult
forall a. a -> StateT EvalState (ExceptT EgisonError RuntimeM) a
forall (m :: * -> *) a. Monad m => a -> m a
return EnvBuildResult
result { ebrTypeEnv = typeEnv' }
  
  -- 5. Pattern Inductive Declarations (from PatternInductiveDecl)
  PatternInductiveDecl String
typeName [String]
typeParams [PatternConstructor]
constructors -> do
    let typeParamVars :: [Type]
typeParamVars = (String -> Type) -> [String] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map (TyVar -> Type
TVar (TyVar -> Type) -> (String -> TyVar) -> String -> Type
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> TyVar
TyVar) [String]
typeParams
        -- Special cases: [a] as TCollection and String as TString
        patternType :: Type
patternType = case (String
typeName, [String]
typeParams) of
                        (String
"[]", [String
param]) -> Type -> Type
TCollection (TyVar -> Type
TVar (String -> TyVar
TyVar String
param))
                        (String
"String", [])  -> Type
TString
                        (String, [String])
_               -> String -> [Type] -> Type
TInductive String
typeName [Type]
typeParamVars
        patternCtorEnv :: PatternConstructorEnv
patternCtorEnv = EnvBuildResult -> PatternConstructorEnv
ebrPatternConstructorEnv EnvBuildResult
result
    
    -- Register each pattern constructor to pattern constructor environment
    PatternConstructorEnv
patternCtorEnv' <- (PatternConstructorEnv
 -> PatternConstructor
 -> StateT
      EvalState (ExceptT EgisonError RuntimeM) PatternConstructorEnv)
-> PatternConstructorEnv
-> [PatternConstructor]
-> StateT
     EvalState (ExceptT EgisonError RuntimeM) PatternConstructorEnv
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM (String
-> [String]
-> Type
-> PatternConstructorEnv
-> PatternConstructor
-> StateT
     EvalState (ExceptT EgisonError RuntimeM) PatternConstructorEnv
registerPatternConstructor String
typeName [String]
typeParams Type
patternType) 
                              PatternConstructorEnv
patternCtorEnv 
                              [PatternConstructor]
constructors
    
    EnvBuildResult -> EvalM EnvBuildResult
forall a. a -> StateT EvalState (ExceptT EgisonError RuntimeM) a
forall (m :: * -> *) a. Monad m => a -> m a
return EnvBuildResult
result { ebrPatternConstructorEnv = patternCtorEnv' }
  
  -- 6. Pattern Function Declarations (from PatternFunctionDecl)
  PatternFunctionDecl String
name [String]
typeParams [(String, TypeExpr)]
params TypeExpr
retType Pattern
_body -> do
    let paramTypes :: [Type]
paramTypes = ((String, TypeExpr) -> Type) -> [(String, TypeExpr)] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map (TypeExpr -> Type
typeExprToType (TypeExpr -> Type)
-> ((String, TypeExpr) -> TypeExpr) -> (String, TypeExpr) -> Type
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (String, TypeExpr) -> TypeExpr
forall a b. (a, b) -> b
snd) [(String, TypeExpr)]
params
        retType' :: Type
retType' = TypeExpr -> Type
typeExprToType TypeExpr
retType
        -- Pattern function type: arg1 -> arg2 -> ... -> retType (without Pattern wrapper)
        patternFuncType :: Type
patternFuncType = (Type -> Type -> Type) -> Type -> [Type] -> Type
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr Type -> Type -> Type
TFun Type
retType' [Type]
paramTypes
        
        -- Quantify over type parameters
        tyVars :: [TyVar]
tyVars = (String -> TyVar) -> [String] -> [TyVar]
forall a b. (a -> b) -> [a] -> [b]
map String -> TyVar
TyVar [String]
typeParams
        typeScheme :: TypeScheme
typeScheme = [TyVar] -> [Constraint] -> Type -> TypeScheme
Types.Forall [TyVar]
tyVars [] Type
patternFuncType
        
        patternEnv :: PatternConstructorEnv
patternEnv = EnvBuildResult -> PatternConstructorEnv
ebrPatternTypeEnv EnvBuildResult
result
        patternEnv' :: PatternConstructorEnv
patternEnv' = String
-> TypeScheme -> PatternConstructorEnv -> PatternConstructorEnv
extendPatternEnv String
name TypeScheme
typeScheme PatternConstructorEnv
patternEnv
    
    EnvBuildResult -> EvalM EnvBuildResult
forall a. a -> StateT EvalState (ExceptT EgisonError RuntimeM) a
forall (m :: * -> *) a. Monad m => a -> m a
return EnvBuildResult
result { ebrPatternTypeEnv = patternEnv' }
  
  -- Other expressions don't contribute to environment building
  Define {} -> EnvBuildResult -> EvalM EnvBuildResult
forall a. a -> StateT EvalState (ExceptT EgisonError RuntimeM) a
forall (m :: * -> *) a. Monad m => a -> m a
return EnvBuildResult
result
  DefineWithType {} -> EnvBuildResult -> EvalM EnvBuildResult
forall a. a -> StateT EvalState (ExceptT EgisonError RuntimeM) a
forall (m :: * -> *) a. Monad m => a -> m a
return EnvBuildResult
result
  Test {} -> EnvBuildResult -> EvalM EnvBuildResult
forall a. a -> StateT EvalState (ExceptT EgisonError RuntimeM) a
forall (m :: * -> *) a. Monad m => a -> m a
return EnvBuildResult
result
  Execute {} -> EnvBuildResult -> EvalM EnvBuildResult
forall a. a -> StateT EvalState (ExceptT EgisonError RuntimeM) a
forall (m :: * -> *) a. Monad m => a -> m a
return EnvBuildResult
result
  LoadFile {} -> EnvBuildResult -> EvalM EnvBuildResult
forall a. a -> StateT EvalState (ExceptT EgisonError RuntimeM) a
forall (m :: * -> *) a. Monad m => a -> m a
return EnvBuildResult
result  -- Should not appear after expandLoads
  InfixDecl {} -> EnvBuildResult -> EvalM EnvBuildResult
forall a. a -> StateT EvalState (ExceptT EgisonError RuntimeM) a
forall (m :: * -> *) a. Monad m => a -> m a
return EnvBuildResult
result
  
  -- 7. Symbol Declarations (from DeclareSymbol)
  DeclareSymbol [String]
names Maybe TypeExpr
mTypeExpr -> do
    let ty :: Type
ty = case Maybe TypeExpr
mTypeExpr of
               Just TypeExpr
texpr -> TypeExpr -> Type
typeExprToType TypeExpr
texpr
               Maybe TypeExpr
Nothing    -> Type
TInt  -- Default to Integer (MathExpr)
        scheme :: TypeScheme
scheme = [TyVar] -> [Constraint] -> Type -> TypeScheme
Forall [] [] Type
ty
        typeEnv :: TypeEnv
typeEnv = EnvBuildResult -> TypeEnv
ebrTypeEnv EnvBuildResult
result
        -- Add each symbol to the type environment
        typeEnv' :: TypeEnv
typeEnv' = (String -> TypeEnv -> TypeEnv) -> TypeEnv -> [String] -> TypeEnv
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\String
name TypeEnv
env -> Var -> TypeScheme -> TypeEnv -> TypeEnv
extendEnv (String -> Var
stringToVar String
name) TypeScheme
scheme TypeEnv
env) TypeEnv
typeEnv [String]
names
    EnvBuildResult -> EvalM EnvBuildResult
forall a. a -> StateT EvalState (ExceptT EgisonError RuntimeM) a
forall (m :: * -> *) a. Monad m => a -> m a
return EnvBuildResult
result { ebrTypeEnv = typeEnv' }

--------------------------------------------------------------------------------
-- Helper Functions
--------------------------------------------------------------------------------

-- | Register a single data constructor
registerConstructor :: String -> [String] -> Type 
                    -> (TypeEnv, ConstructorEnv) -> InductiveConstructor 
                    -> EvalM (TypeEnv, ConstructorEnv)
registerConstructor :: String
-> [String]
-> Type
-> (TypeEnv, ConstructorEnv)
-> InductiveConstructor
-> StateT
     EvalState (ExceptT EgisonError RuntimeM) (TypeEnv, ConstructorEnv)
registerConstructor String
typeName [String]
typeParams Type
resultType (TypeEnv
typeEnv, ConstructorEnv
ctorEnv) 
                    (InductiveConstructor String
ctorName [TypeExpr]
argTypeExprs) = do
  let argTypes :: [Type]
argTypes = (TypeExpr -> Type) -> [TypeExpr] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map TypeExpr -> Type
typeExprToType [TypeExpr]
argTypeExprs
      
      -- Constructor type: argTypes -> resultType
      constructorType :: Type
constructorType = (Type -> Type -> Type) -> Type -> [Type] -> Type
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr Type -> Type -> Type
TFun Type
resultType [Type]
argTypes
      
      -- Quantify over type parameters
      tyVars :: [TyVar]
tyVars = (String -> TyVar) -> [String] -> [TyVar]
forall a b. (a -> b) -> [a] -> [b]
map String -> TyVar
TyVar [String]
typeParams
      typeScheme :: TypeScheme
typeScheme = [TyVar] -> [Constraint] -> Type -> TypeScheme
Types.Forall [TyVar]
tyVars [] Type
constructorType
      
      -- Add to type environment
      typeEnv' :: TypeEnv
typeEnv' = Var -> TypeScheme -> TypeEnv -> TypeEnv
extendEnv (String -> Var
stringToVar String
ctorName) TypeScheme
typeScheme TypeEnv
typeEnv
      
      -- Add to constructor environment (for pattern matching and evaluation)
      ctorInfo :: ConstructorInfo
ctorInfo = ConstructorInfo
        { ctorTypeName :: String
ctorTypeName = String
typeName
        , ctorArgTypes :: [Type]
ctorArgTypes = [Type]
argTypes
        , ctorTypeParams :: [String]
ctorTypeParams = [String]
typeParams
        }
      ctorEnv' :: ConstructorEnv
ctorEnv' = String -> ConstructorInfo -> ConstructorEnv -> ConstructorEnv
forall k v. Hashable k => k -> v -> HashMap k v -> HashMap k v
HashMap.insert String
ctorName ConstructorInfo
ctorInfo ConstructorEnv
ctorEnv
  
  (TypeEnv, ConstructorEnv)
-> StateT
     EvalState (ExceptT EgisonError RuntimeM) (TypeEnv, ConstructorEnv)
forall a. a -> StateT EvalState (ExceptT EgisonError RuntimeM) a
forall (m :: * -> *) a. Monad m => a -> m a
return (TypeEnv
typeEnv', ConstructorEnv
ctorEnv')

-- | Register a class method to the type environment
registerClassMethod :: TyVar -> String -> TypeEnv -> ClassMethod -> TypeEnv
registerClassMethod :: TyVar -> String -> TypeEnv -> ClassMethod -> TypeEnv
registerClassMethod TyVar
tyVar String
className TypeEnv
typeEnv (ClassMethod String
methName [TypedParam]
params TypeExpr
retType Maybe Expr
_defaultImpl) =
  let paramTypes :: [Type]
paramTypes = (TypedParam -> Type) -> [TypedParam] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map TypedParam -> Type
typedParamToType [TypedParam]
params
      methodType :: Type
methodType = (Type -> Type -> Type) -> Type -> [Type] -> Type
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr Type -> Type -> Type
TFun (TypeExpr -> Type
typeExprToType TypeExpr
retType) [Type]
paramTypes
      
      -- Method has constrained type: ClassName a => methodType
      constraint :: Constraint
constraint = String -> Type -> Constraint
Types.Constraint String
className (TyVar -> Type
TVar TyVar
tyVar)
      typeScheme :: TypeScheme
typeScheme = [TyVar] -> [Constraint] -> Type -> TypeScheme
Types.Forall [TyVar
tyVar] [Constraint
constraint] Type
methodType
  in
    Var -> TypeScheme -> TypeEnv -> TypeEnv
extendEnv (String -> Var
stringToVar String
methName) TypeScheme
typeScheme TypeEnv
typeEnv

-- | Register type signatures for instance methods (generated during desugaring)
-- This prevents "Unbound variable" warnings during type inference
registerInstanceMethods :: String -> Type -> [Constraint] -> [InstanceMethod] -> ClassEnv -> TypeEnv -> TypeEnv
registerInstanceMethods :: String
-> Type
-> [Constraint]
-> [InstanceMethod]
-> ClassEnv
-> TypeEnv
-> TypeEnv
registerInstanceMethods String
className Type
instType [Constraint]
instConstraints [InstanceMethod]
methods ClassEnv
classEnv TypeEnv
typeEnv =
  case String -> ClassEnv -> Maybe ClassInfo
lookupClass String
className ClassEnv
classEnv of
    Maybe ClassInfo
Nothing -> TypeEnv
typeEnv  -- Class not found, skip
    Just ClassInfo
classInfo -> 
      -- Register each instance method
      let typeEnv' :: TypeEnv
typeEnv' = (InstanceMethod -> TypeEnv -> TypeEnv)
-> TypeEnv -> [InstanceMethod] -> TypeEnv
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (String
-> Type
-> [Constraint]
-> ClassInfo
-> InstanceMethod
-> TypeEnv
-> TypeEnv
registerInstanceMethod String
className Type
instType [Constraint]
instConstraints ClassInfo
classInfo) TypeEnv
typeEnv [InstanceMethod]
methods
      
          -- Also register the dictionary itself
          -- e.g., eqCollection : {Eq a} Hash String ([a] -> [a] -> Bool)
          typeName' :: String
typeName' = Type -> String
Types.typeConstructorName Type
instType
          dictName :: String
dictName = ShowS
lowerFirst String
className String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
typeName'
          
          -- Build dictionary type: Hash String (method type)
          -- All methods should have the same general shape, so we use the first one
          dictValueType :: Type
dictValueType = case [InstanceMethod]
methods of
            [] -> Type
TAny
            [InstanceMethod]
_ -> case String -> [(String, Type)] -> Maybe Type
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup (InstanceMethod -> String
instanceMethodName ([InstanceMethod] -> InstanceMethod
forall a. HasCallStack => [a] -> a
head [InstanceMethod]
methods)) (ClassInfo -> [(String, Type)]
Types.classMethods ClassInfo
classInfo) of
              Maybe Type
Nothing -> Type
TAny
              Just Type
methodType ->
                let tyVar :: TyVar
tyVar = ClassInfo -> TyVar
Types.classParam ClassInfo
classInfo
                    substitutedType :: Type
substitutedType = TyVar -> Type -> Type -> Type
substituteTypeVar TyVar
tyVar Type
instType Type
methodType
                in Type
substitutedType
          
          dictType :: Type
dictType = Type -> Type -> Type
THash Type
TString Type
dictValueType
          freeVars :: [TyVar]
freeVars = Set TyVar -> [TyVar]
forall a. Set a -> [a]
Set.toList (Type -> Set TyVar
freeTyVars Type
dictType)
          dictScheme :: TypeScheme
dictScheme = [TyVar] -> [Constraint] -> Type -> TypeScheme
Types.Forall [TyVar]
freeVars [Constraint]
instConstraints Type
dictType
      in
        Var -> TypeScheme -> TypeEnv -> TypeEnv
extendEnv (String -> Var
stringToVar String
dictName) TypeScheme
dictScheme TypeEnv
typeEnv'
  where
    instanceMethodName :: InstanceMethod -> String
    instanceMethodName :: InstanceMethod -> String
instanceMethodName (InstanceMethod String
name [String]
_ Expr
_) = String
name
    
    registerInstanceMethod :: String -> Type -> [Constraint] -> Types.ClassInfo -> InstanceMethod -> TypeEnv -> TypeEnv
    registerInstanceMethod :: String
-> Type
-> [Constraint]
-> ClassInfo
-> InstanceMethod
-> TypeEnv
-> TypeEnv
registerInstanceMethod String
clsName Type
instTy [Constraint]
constraints ClassInfo
classInfo (InstanceMethod String
methName [String]
_params Expr
_body) TypeEnv
env =
      -- Find the method in the class definition
      case String -> [(String, Type)] -> Maybe Type
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup String
methName (ClassInfo -> [(String, Type)]
Types.classMethods ClassInfo
classInfo) of
        Maybe Type
Nothing -> TypeEnv
env  -- Method not in class definition, skip
        Just Type
methodType -> 
          -- Substitute type variable with instance type
          let tyVar :: TyVar
tyVar = ClassInfo -> TyVar
Types.classParam ClassInfo
classInfo
              substitutedType :: Type
substitutedType = TyVar -> Type -> Type -> Type
substituteTypeVar TyVar
tyVar Type
instTy Type
methodType
              
              -- Generate method name using type constructor name only (no type parameters)
              -- e.g., "eqCollectionEq" not "eqCollectionaEq"
              typeName' :: String
typeName' = Type -> String
Types.typeConstructorName Type
instTy
              sanitizedName :: String
sanitizedName = ShowS
sanitizeMethodName String
methName
              generatedMethodName :: String
generatedMethodName = ShowS
lowerFirst String
clsName String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
typeName' String -> ShowS
forall a. [a] -> [a] -> [a]
++ ShowS
capitalizeFirst String
sanitizedName
              
              -- Extract free type variables from the substituted type
              freeVars :: [TyVar]
freeVars = Set TyVar -> [TyVar]
forall a. Set a -> [a]
Set.toList (Type -> Set TyVar
freeTyVars Type
substitutedType)
              
              -- Create type scheme with constraints from the instance context
              -- e.g., {Eq a} [a] -> [a] -> Bool for instance {Eq a} Eq [a]
              typeScheme :: TypeScheme
typeScheme = [TyVar] -> [Constraint] -> Type -> TypeScheme
Types.Forall [TyVar]
freeVars [Constraint]
constraints Type
substitutedType
          in
            Var -> TypeScheme -> TypeEnv -> TypeEnv
extendEnv (String -> Var
stringToVar String
generatedMethodName) TypeScheme
typeScheme TypeEnv
env
    
    -- Substitute type variable with concrete type in a type expression
    substituteTypeVar :: TyVar -> Type -> Type -> Type
    substituteTypeVar :: TyVar -> Type -> Type -> Type
substituteTypeVar TyVar
oldVar Type
newType = Type -> Type
go
      where
        go :: Type -> Type
go Type
TInt = Type
TInt
        go Type
TFloat = Type
TFloat
        go Type
TBool = Type
TBool
        go Type
TChar = Type
TChar
        go Type
TString = Type
TString
        go (TVar TyVar
v) | TyVar
v TyVar -> TyVar -> Bool
forall a. Eq a => a -> a -> Bool
== TyVar
oldVar = Type
newType
                    | Bool
otherwise = TyVar -> Type
TVar TyVar
v
        go (TTuple [Type]
ts) = [Type] -> Type
TTuple ((Type -> Type) -> [Type] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map Type -> Type
go [Type]
ts)
        go (TCollection Type
t) = Type -> Type
TCollection (Type -> Type
go Type
t)
        go (TInductive String
name [Type]
ts) = String -> [Type] -> Type
TInductive String
name ((Type -> Type) -> [Type] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map Type -> Type
go [Type]
ts)
        go (TTensor Type
t) = Type -> Type
TTensor (Type -> Type
go Type
t)
        go (THash Type
k Type
v) = Type -> Type -> Type
THash (Type -> Type
go Type
k) (Type -> Type
go Type
v)
        go (TMatcher Type
t) = Type -> Type
TMatcher (Type -> Type
go Type
t)
        go (TFun Type
t1 Type
t2) = Type -> Type -> Type
TFun (Type -> Type
go Type
t1) (Type -> Type
go Type
t2)
        go (TIO Type
t) = Type -> Type
TIO (Type -> Type
go Type
t)
        go (TIORef Type
t) = Type -> Type
TIORef (Type -> Type
go Type
t)
        go Type
TAny = Type
TAny

-- | Extract method name from ClassMethod
extractMethodName :: ClassMethod -> String
extractMethodName :: ClassMethod -> String
extractMethodName (ClassMethod String
name [TypedParam]
_ TypeExpr
_ Maybe Expr
_) = String
name

-- | Extract method name and type from ClassMethod
extractMethodWithType :: ClassMethod -> (String, Type)
extractMethodWithType :: ClassMethod -> (String, Type)
extractMethodWithType (ClassMethod String
name [TypedParam]
params TypeExpr
retType Maybe Expr
_) =
  let paramTypes :: [Type]
paramTypes = (TypedParam -> Type) -> [TypedParam] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map TypedParam -> Type
typedParamToType [TypedParam]
params
      methodType :: Type
methodType = (Type -> Type -> Type) -> Type -> [Type] -> Type
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr Type -> Type -> Type
TFun (TypeExpr -> Type
typeExprToType TypeExpr
retType) [Type]
paramTypes
  in (String
name, Type
methodType)

-- | Extract class name from ConstraintExpr
extractConstraintName :: ConstraintExpr -> String
extractConstraintName :: ConstraintExpr -> String
extractConstraintName (ConstraintExpr String
clsName [TypeExpr]
_) = String
clsName

-- | Convert ConstraintExpr to internal Constraint
constraintToInternal :: ConstraintExpr -> Types.Constraint
constraintToInternal :: ConstraintExpr -> Constraint
constraintToInternal (ConstraintExpr String
clsName [TypeExpr]
tyExprs) =
  String -> Type -> Constraint
Types.Constraint String
clsName (case [TypeExpr]
tyExprs of 
    [] -> Type
TAny
    (TypeExpr
t:[TypeExpr]
_) -> TypeExpr -> Type
typeExprToType TypeExpr
t)

-- | Register a single pattern constructor
registerPatternConstructor :: String -> [String] -> Type 
                           -> PatternConstructorEnv -> PatternConstructor 
                           -> EvalM PatternConstructorEnv
registerPatternConstructor :: String
-> [String]
-> Type
-> PatternConstructorEnv
-> PatternConstructor
-> StateT
     EvalState (ExceptT EgisonError RuntimeM) PatternConstructorEnv
registerPatternConstructor String
_typeName [String]
typeParams Type
resultType PatternConstructorEnv
patternCtorEnv 
                          (PatternConstructor String
ctorName [TypeExpr]
argTypeExprs) = do
  let argTypes :: [Type]
argTypes = (TypeExpr -> Type) -> [TypeExpr] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map TypeExpr -> Type
typeExprToType [TypeExpr]
argTypeExprs
      
      -- Pattern constructor type: arg1 -> arg2 -> ... -> resultType (without Pattern wrapper)
      patternCtorType :: Type
patternCtorType = (Type -> Type -> Type) -> Type -> [Type] -> Type
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr Type -> Type -> Type
TFun Type
resultType [Type]
argTypes
      
      -- Quantify over type parameters
      tyVars :: [TyVar]
tyVars = (String -> TyVar) -> [String] -> [TyVar]
forall a b. (a -> b) -> [a] -> [b]
map String -> TyVar
TyVar [String]
typeParams
      typeScheme :: TypeScheme
typeScheme = [TyVar] -> [Constraint] -> Type -> TypeScheme
Types.Forall [TyVar]
tyVars [] Type
patternCtorType
      
      -- Add to pattern constructor environment (same format as PatternTypeEnv)
      patternCtorEnv' :: PatternConstructorEnv
patternCtorEnv' = String
-> TypeScheme -> PatternConstructorEnv -> PatternConstructorEnv
extendPatternEnv String
ctorName TypeScheme
typeScheme PatternConstructorEnv
patternCtorEnv
  
  PatternConstructorEnv
-> StateT
     EvalState (ExceptT EgisonError RuntimeM) PatternConstructorEnv
forall a. a -> StateT EvalState (ExceptT EgisonError RuntimeM) a
forall (m :: * -> *) a. Monad m => a -> m a
return PatternConstructorEnv
patternCtorEnv'

-- | Convert TypedParam to Type
typedParamToType :: TypedParam -> Type
typedParamToType :: TypedParam -> Type
typedParamToType (TPVar String
_ TypeExpr
ty) = TypeExpr -> Type
typeExprToType TypeExpr
ty
typedParamToType (TPInvertedVar String
_ TypeExpr
ty) = TypeExpr -> Type
typeExprToType TypeExpr
ty
typedParamToType (TPTuple [TypedParam]
elems) = [Type] -> Type
TTuple ((TypedParam -> Type) -> [TypedParam] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map TypedParam -> Type
typedParamToType [TypedParam]
elems)
typedParamToType (TPWildcard TypeExpr
ty) = TypeExpr -> Type
typeExprToType TypeExpr
ty
typedParamToType (TPUntypedVar String
_) = TyVar -> Type
TVar (String -> TyVar
TyVar String
"a")  -- Will be inferred
typedParamToType TypedParam
TPUntypedWildcard = TyVar -> Type
TVar (String -> TyVar
TyVar String
"a")  -- Will be inferred