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

This module provides type environment for the Egison type system.
-}

module Language.Egison.Type.Env
  ( TypeEnv(..)
  , emptyEnv
  , extendEnv
  , extendEnvMany
  , lookupEnv
  , removeFromEnv
  , envToList
  , freeVarsInEnv
  , generalize
  , instantiate
  -- * Class environment
  , ClassEnv(..)
  , ClassInfo(..)
  , InstanceInfo(..)
  , emptyClassEnv
  , addClass
  , addInstance
  , lookupClass
  , lookupInstances
  , classEnvToList
  , mergeClassEnv
  -- * Pattern type environment
  , PatternTypeEnv(..)
  , emptyPatternEnv
  , extendPatternEnv
  , lookupPatternEnv
  , patternEnvToList
  ) where

import           Data.List                  (sortOn)
import           Data.Map.Strict            (Map)
import qualified Data.Map.Strict            as Map
import           Data.Set                   (Set)
import qualified Data.Set                   as Set

import           Language.Egison.IExpr      (Var(..), Index(..))
import           Language.Egison.VarEntry   (VarEntry(..))
import           Language.Egison.Type.Types (TyVar (..), Type (..), TypeScheme (..),
                                             Constraint(..), ClassInfo(..), InstanceInfo(..),
                                             freeTyVars, freshTyVar)

-- | Type environment: uses same data structure as evaluation environment
-- Maps base variable names to all bindings with that name
-- VarEntry list is sorted by index length (shortest first) for efficient prefix matching
newtype TypeEnv = TypeEnv { TypeEnv -> Map String [VarEntry TypeScheme]
unTypeEnv :: Map String [VarEntry TypeScheme] }
  deriving (TypeEnv -> TypeEnv -> Bool
(TypeEnv -> TypeEnv -> Bool)
-> (TypeEnv -> TypeEnv -> Bool) -> Eq TypeEnv
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: TypeEnv -> TypeEnv -> Bool
== :: TypeEnv -> TypeEnv -> Bool
$c/= :: TypeEnv -> TypeEnv -> Bool
/= :: TypeEnv -> TypeEnv -> Bool
Eq, Int -> TypeEnv -> ShowS
[TypeEnv] -> ShowS
TypeEnv -> String
(Int -> TypeEnv -> ShowS)
-> (TypeEnv -> String) -> ([TypeEnv] -> ShowS) -> Show TypeEnv
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> TypeEnv -> ShowS
showsPrec :: Int -> TypeEnv -> ShowS
$cshow :: TypeEnv -> String
show :: TypeEnv -> String
$cshowList :: [TypeEnv] -> ShowS
showList :: [TypeEnv] -> ShowS
Show)

-- | Pattern type environment: maps pattern function names to type schemes
-- This is separate from the value type environment
newtype PatternTypeEnv = PatternTypeEnv { PatternTypeEnv -> Map String TypeScheme
unPatternTypeEnv :: Map String TypeScheme }
  deriving (PatternTypeEnv -> PatternTypeEnv -> Bool
(PatternTypeEnv -> PatternTypeEnv -> Bool)
-> (PatternTypeEnv -> PatternTypeEnv -> Bool) -> Eq PatternTypeEnv
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: PatternTypeEnv -> PatternTypeEnv -> Bool
== :: PatternTypeEnv -> PatternTypeEnv -> Bool
$c/= :: PatternTypeEnv -> PatternTypeEnv -> Bool
/= :: PatternTypeEnv -> PatternTypeEnv -> Bool
Eq, Int -> PatternTypeEnv -> ShowS
[PatternTypeEnv] -> ShowS
PatternTypeEnv -> String
(Int -> PatternTypeEnv -> ShowS)
-> (PatternTypeEnv -> String)
-> ([PatternTypeEnv] -> ShowS)
-> Show PatternTypeEnv
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> PatternTypeEnv -> ShowS
showsPrec :: Int -> PatternTypeEnv -> ShowS
$cshow :: PatternTypeEnv -> String
show :: PatternTypeEnv -> String
$cshowList :: [PatternTypeEnv] -> ShowS
showList :: [PatternTypeEnv] -> ShowS
Show)

-- | Empty type environment
emptyEnv :: TypeEnv
emptyEnv :: TypeEnv
emptyEnv = Map String [VarEntry TypeScheme] -> TypeEnv
TypeEnv Map String [VarEntry TypeScheme]
forall k a. Map k a
Map.empty

-- | Extend the environment with a new binding
extendEnv :: Var -> TypeScheme -> TypeEnv -> TypeEnv
extendEnv :: Var -> TypeScheme -> TypeEnv -> TypeEnv
extendEnv (Var String
name [Index (Maybe Var)]
indices) TypeScheme
scheme (TypeEnv Map String [VarEntry TypeScheme]
env) =
  let entry :: VarEntry TypeScheme
entry = [Index (Maybe Var)] -> TypeScheme -> VarEntry TypeScheme
forall a. [Index (Maybe Var)] -> a -> VarEntry a
VarEntry [Index (Maybe Var)]
indices TypeScheme
scheme
      newEntries :: [VarEntry TypeScheme]
newEntries = case String
-> Map String [VarEntry TypeScheme] -> Maybe [VarEntry TypeScheme]
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup String
name Map String [VarEntry TypeScheme]
env of
        Maybe [VarEntry TypeScheme]
Nothing -> [VarEntry TypeScheme
entry]
        Just [VarEntry TypeScheme]
existingEntries -> (VarEntry TypeScheme -> Int)
-> [VarEntry TypeScheme] -> [VarEntry TypeScheme]
forall b a. Ord b => (a -> b) -> [a] -> [a]
sortOn ([Index (Maybe Var)] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([Index (Maybe Var)] -> Int)
-> (VarEntry TypeScheme -> [Index (Maybe Var)])
-> VarEntry TypeScheme
-> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VarEntry TypeScheme -> [Index (Maybe Var)]
forall a. VarEntry a -> [Index (Maybe Var)]
veIndices) (VarEntry TypeScheme
entry VarEntry TypeScheme
-> [VarEntry TypeScheme] -> [VarEntry TypeScheme]
forall a. a -> [a] -> [a]
: [VarEntry TypeScheme]
existingEntries)
  in Map String [VarEntry TypeScheme] -> TypeEnv
TypeEnv (Map String [VarEntry TypeScheme] -> TypeEnv)
-> Map String [VarEntry TypeScheme] -> TypeEnv
forall a b. (a -> b) -> a -> b
$ String
-> [VarEntry TypeScheme]
-> Map String [VarEntry TypeScheme]
-> Map String [VarEntry TypeScheme]
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert String
name [VarEntry TypeScheme]
newEntries Map String [VarEntry TypeScheme]
env

-- | Extend the environment with multiple bindings
extendEnvMany :: [(Var, TypeScheme)] -> TypeEnv -> TypeEnv
extendEnvMany :: [(Var, TypeScheme)] -> TypeEnv -> TypeEnv
extendEnvMany [(Var, TypeScheme)]
bindings TypeEnv
env = ((Var, TypeScheme) -> TypeEnv -> TypeEnv)
-> TypeEnv -> [(Var, TypeScheme)] -> TypeEnv
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr ((Var -> TypeScheme -> TypeEnv -> TypeEnv)
-> (Var, TypeScheme) -> TypeEnv -> TypeEnv
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry Var -> TypeScheme -> TypeEnv -> TypeEnv
extendEnv) TypeEnv
env [(Var, TypeScheme)]
bindings

-- | Look up a variable in the environment
-- Search algorithm (same as refVar in Data.hs):
--   1. Try exact match
--   2. Try prefix match (find longer indices and auto-complete with #)
--   3. Try suffix removal (find shorter indices, pick longest match)
-- No recursion is used; all matching is done in a single pass to avoid infinite loops.
lookupEnv :: Var -> TypeEnv -> Maybe TypeScheme
lookupEnv :: Var -> TypeEnv -> Maybe TypeScheme
lookupEnv (Var String
name [Index (Maybe Var)]
targetIndices) (TypeEnv Map String [VarEntry TypeScheme]
env) =
  case String
-> Map String [VarEntry TypeScheme] -> Maybe [VarEntry TypeScheme]
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup String
name Map String [VarEntry TypeScheme]
env of
    Maybe [VarEntry TypeScheme]
Nothing -> Maybe TypeScheme
forall a. Maybe a
Nothing
    Just [VarEntry TypeScheme]
entries ->
      -- 1. Try exact match first
      case [Index (Maybe Var)] -> [VarEntry TypeScheme] -> Maybe TypeScheme
findExactMatch [Index (Maybe Var)]
targetIndices [VarEntry TypeScheme]
entries of
        Just TypeScheme
scheme -> TypeScheme -> Maybe TypeScheme
forall a. a -> Maybe a
Just TypeScheme
scheme
        Maybe TypeScheme
Nothing ->
          -- 2. Try prefix matching (e_a matches e_i_j)
          case [Index (Maybe Var)] -> [VarEntry TypeScheme] -> Maybe TypeScheme
findPrefixMatch [Index (Maybe Var)]
targetIndices [VarEntry TypeScheme]
entries of
            Just TypeScheme
scheme -> TypeScheme -> Maybe TypeScheme
forall a. a -> Maybe a
Just TypeScheme
scheme
            Maybe TypeScheme
Nothing ->
              -- 3. Try suffix removal (e_i_j_k matches e_i_j, pick longest)
              [Index (Maybe Var)] -> [VarEntry TypeScheme] -> Maybe TypeScheme
findSuffixMatch [Index (Maybe Var)]
targetIndices [VarEntry TypeScheme]
entries
  where
    -- Exact match: same length and same indices
    findExactMatch :: [Index (Maybe Var)] -> [VarEntry TypeScheme] -> Maybe TypeScheme
    findExactMatch :: [Index (Maybe Var)] -> [VarEntry TypeScheme] -> Maybe TypeScheme
findExactMatch [Index (Maybe Var)]
indices [VarEntry TypeScheme]
entries =
      case [VarEntry TypeScheme -> TypeScheme
forall a. VarEntry a -> a
veValue VarEntry TypeScheme
e | VarEntry TypeScheme
e <- [VarEntry TypeScheme]
entries, VarEntry TypeScheme -> [Index (Maybe Var)]
forall a. VarEntry a -> [Index (Maybe Var)]
veIndices VarEntry TypeScheme
e [Index (Maybe Var)] -> [Index (Maybe Var)] -> Bool
forall a. Eq a => a -> a -> Bool
== [Index (Maybe Var)]
indices] of
        (TypeScheme
scheme:[TypeScheme]
_) -> TypeScheme -> Maybe TypeScheme
forall a. a -> Maybe a
Just TypeScheme
scheme
        [] -> Maybe TypeScheme
forall a. Maybe a
Nothing
    
    -- Prefix matching: find shortest entry where target indices are a prefix
    -- Example: target [a] matches [i, j] in e_i_j (shortest match)
    findPrefixMatch :: [Index (Maybe Var)] -> [VarEntry TypeScheme] -> Maybe TypeScheme
    findPrefixMatch :: [Index (Maybe Var)] -> [VarEntry TypeScheme] -> Maybe TypeScheme
findPrefixMatch [Index (Maybe Var)]
indices [VarEntry TypeScheme]
entries =
      -- entries are sorted by index length (ascending), so first match is shortest
      case [VarEntry TypeScheme -> TypeScheme
forall a. VarEntry a -> a
veValue VarEntry TypeScheme
e | VarEntry TypeScheme
e <- [VarEntry TypeScheme]
entries, [Index (Maybe Var)] -> [Index (Maybe Var)] -> Bool
isPrefixOfIndices [Index (Maybe Var)]
indices (VarEntry TypeScheme -> [Index (Maybe Var)]
forall a. VarEntry a -> [Index (Maybe Var)]
veIndices VarEntry TypeScheme
e)] of
        (TypeScheme
scheme:[TypeScheme]
_) -> TypeScheme -> Maybe TypeScheme
forall a. a -> Maybe a
Just TypeScheme
scheme
        [] -> Maybe TypeScheme
forall a. Maybe a
Nothing
    
    -- Suffix removal: find longest entry where stored indices are a prefix of target
    -- Example: target [i,j,k] matches e_i_j (stored [i,j]); prefer e_i_j over e_i
    -- Single pass, no recursion - safe from infinite loops
    findSuffixMatch :: [Index (Maybe Var)] -> [VarEntry TypeScheme] -> Maybe TypeScheme
    findSuffixMatch :: [Index (Maybe Var)] -> [VarEntry TypeScheme] -> Maybe TypeScheme
findSuffixMatch [Index (Maybe Var)]
targetIndices [VarEntry TypeScheme]
entries =
      let suffixMatches :: [VarEntry TypeScheme]
suffixMatches = [VarEntry TypeScheme
e | VarEntry TypeScheme
e <- [VarEntry TypeScheme]
entries, [Index (Maybe Var)] -> [Index (Maybe Var)] -> Bool
storedIsPrefixOfTarget (VarEntry TypeScheme -> [Index (Maybe Var)]
forall a. VarEntry a -> [Index (Maybe Var)]
veIndices VarEntry TypeScheme
e) [Index (Maybe Var)]
targetIndices]
      in case [VarEntry TypeScheme] -> [VarEntry TypeScheme]
sortByIndexLengthDesc [VarEntry TypeScheme]
suffixMatches of
        (VarEntry TypeScheme
e:[VarEntry TypeScheme]
_) -> TypeScheme -> Maybe TypeScheme
forall a. a -> Maybe a
Just (VarEntry TypeScheme -> TypeScheme
forall a. VarEntry a -> a
veValue VarEntry TypeScheme
e)
        [] -> Maybe TypeScheme
forall a. Maybe a
Nothing
    
    -- stored is prefix of target: stored has fewer indices, first part of target matches
    storedIsPrefixOfTarget :: [Index (Maybe Var)] -> [Index (Maybe Var)] -> Bool
    storedIsPrefixOfTarget :: [Index (Maybe Var)] -> [Index (Maybe Var)] -> Bool
storedIsPrefixOfTarget [Index (Maybe Var)]
stored [Index (Maybe Var)]
target =
      Bool -> Bool
not ([Index (Maybe Var)] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Index (Maybe Var)]
target) Bool -> Bool -> Bool
&&
      [Index (Maybe Var)] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Index (Maybe Var)]
stored Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< [Index (Maybe Var)] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Index (Maybe Var)]
target Bool -> Bool -> Bool
&&
      [Index (Maybe Var)]
stored [Index (Maybe Var)] -> [Index (Maybe Var)] -> Bool
forall a. Eq a => a -> a -> Bool
== Int -> [Index (Maybe Var)] -> [Index (Maybe Var)]
forall a. Int -> [a] -> [a]
take ([Index (Maybe Var)] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Index (Maybe Var)]
stored) [Index (Maybe Var)]
target
    
    sortByIndexLengthDesc :: [VarEntry TypeScheme] -> [VarEntry TypeScheme]
    sortByIndexLengthDesc :: [VarEntry TypeScheme] -> [VarEntry TypeScheme]
sortByIndexLengthDesc = [VarEntry TypeScheme] -> [VarEntry TypeScheme]
forall a. [a] -> [a]
reverse ([VarEntry TypeScheme] -> [VarEntry TypeScheme])
-> ([VarEntry TypeScheme] -> [VarEntry TypeScheme])
-> [VarEntry TypeScheme]
-> [VarEntry TypeScheme]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (VarEntry TypeScheme -> Int)
-> [VarEntry TypeScheme] -> [VarEntry TypeScheme]
forall b a. Ord b => (a -> b) -> [a] -> [a]
sortOn ([Index (Maybe Var)] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([Index (Maybe Var)] -> Int)
-> (VarEntry TypeScheme -> [Index (Maybe Var)])
-> VarEntry TypeScheme
-> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. VarEntry TypeScheme -> [Index (Maybe Var)]
forall a. VarEntry a -> [Index (Maybe Var)]
veIndices)
    
    -- Check if target is a prefix of candidate (for prefix matching)
    -- Example: [a] is prefix of [i, j]
    -- IMPORTANT: target must be non-empty to avoid matching everything
    isPrefixOfIndices :: [Index (Maybe Var)] -> [Index (Maybe Var)] -> Bool
    isPrefixOfIndices :: [Index (Maybe Var)] -> [Index (Maybe Var)] -> Bool
isPrefixOfIndices [Index (Maybe Var)]
target [Index (Maybe Var)]
candidate =
      Bool -> Bool
not ([Index (Maybe Var)] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Index (Maybe Var)]
target) Bool -> Bool -> Bool
&&
      [Index (Maybe Var)] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Index (Maybe Var)]
target Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< [Index (Maybe Var)] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Index (Maybe Var)]
candidate Bool -> Bool -> Bool
&&
      [Index (Maybe Var)]
target [Index (Maybe Var)] -> [Index (Maybe Var)] -> Bool
forall a. Eq a => a -> a -> Bool
== Int -> [Index (Maybe Var)] -> [Index (Maybe Var)]
forall a. Int -> [a] -> [a]
take ([Index (Maybe Var)] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Index (Maybe Var)]
target) [Index (Maybe Var)]
candidate

-- | Remove a variable from the environment
removeFromEnv :: Var -> TypeEnv -> TypeEnv
removeFromEnv :: Var -> TypeEnv -> TypeEnv
removeFromEnv (Var String
name [Index (Maybe Var)]
indices) (TypeEnv Map String [VarEntry TypeScheme]
env) =
  case String
-> Map String [VarEntry TypeScheme] -> Maybe [VarEntry TypeScheme]
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup String
name Map String [VarEntry TypeScheme]
env of
    Maybe [VarEntry TypeScheme]
Nothing -> Map String [VarEntry TypeScheme] -> TypeEnv
TypeEnv Map String [VarEntry TypeScheme]
env
    Just [VarEntry TypeScheme]
entries ->
      let newEntries :: [VarEntry TypeScheme]
newEntries = [VarEntry TypeScheme
e | VarEntry TypeScheme
e <- [VarEntry TypeScheme]
entries, VarEntry TypeScheme -> [Index (Maybe Var)]
forall a. VarEntry a -> [Index (Maybe Var)]
veIndices VarEntry TypeScheme
e [Index (Maybe Var)] -> [Index (Maybe Var)] -> Bool
forall a. Eq a => a -> a -> Bool
/= [Index (Maybe Var)]
indices]
      in if [VarEntry TypeScheme] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [VarEntry TypeScheme]
newEntries
         then Map String [VarEntry TypeScheme] -> TypeEnv
TypeEnv (Map String [VarEntry TypeScheme] -> TypeEnv)
-> Map String [VarEntry TypeScheme] -> TypeEnv
forall a b. (a -> b) -> a -> b
$ String
-> Map String [VarEntry TypeScheme]
-> Map String [VarEntry TypeScheme]
forall k a. Ord k => k -> Map k a -> Map k a
Map.delete String
name Map String [VarEntry TypeScheme]
env
         else Map String [VarEntry TypeScheme] -> TypeEnv
TypeEnv (Map String [VarEntry TypeScheme] -> TypeEnv)
-> Map String [VarEntry TypeScheme] -> TypeEnv
forall a b. (a -> b) -> a -> b
$ String
-> [VarEntry TypeScheme]
-> Map String [VarEntry TypeScheme]
-> Map String [VarEntry TypeScheme]
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert String
name [VarEntry TypeScheme]
newEntries Map String [VarEntry TypeScheme]
env

-- | Convert environment to list
envToList :: TypeEnv -> [(Var, TypeScheme)]
envToList :: TypeEnv -> [(Var, TypeScheme)]
envToList (TypeEnv Map String [VarEntry TypeScheme]
env) =
  [ (String -> [Index (Maybe Var)] -> Var
Var String
name (VarEntry TypeScheme -> [Index (Maybe Var)]
forall a. VarEntry a -> [Index (Maybe Var)]
veIndices VarEntry TypeScheme
entry), VarEntry TypeScheme -> TypeScheme
forall a. VarEntry a -> a
veValue VarEntry TypeScheme
entry)
  | (String
name, [VarEntry TypeScheme]
entries) <- Map String [VarEntry TypeScheme]
-> [(String, [VarEntry TypeScheme])]
forall k a. Map k a -> [(k, a)]
Map.toList Map String [VarEntry TypeScheme]
env
  , VarEntry TypeScheme
entry <- [VarEntry TypeScheme]
entries
  ]

-- | Get free type variables in the environment
freeVarsInEnv :: TypeEnv -> Set TyVar
freeVarsInEnv :: TypeEnv -> Set TyVar
freeVarsInEnv (TypeEnv Map String [VarEntry TypeScheme]
env) = 
  [Set TyVar] -> Set TyVar
forall (f :: * -> *) a. (Foldable f, Ord a) => f (Set a) -> Set a
Set.unions ([Set TyVar] -> Set TyVar) -> [Set TyVar] -> Set TyVar
forall a b. (a -> b) -> a -> b
$ (VarEntry TypeScheme -> Set TyVar)
-> [VarEntry TypeScheme] -> [Set TyVar]
forall a b. (a -> b) -> [a] -> [b]
map VarEntry TypeScheme -> Set TyVar
freeVarsInScheme ([VarEntry TypeScheme] -> [Set TyVar])
-> [VarEntry TypeScheme] -> [Set TyVar]
forall a b. (a -> b) -> a -> b
$ [[VarEntry TypeScheme]] -> [VarEntry TypeScheme]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[VarEntry TypeScheme]] -> [VarEntry TypeScheme])
-> [[VarEntry TypeScheme]] -> [VarEntry TypeScheme]
forall a b. (a -> b) -> a -> b
$ Map String [VarEntry TypeScheme] -> [[VarEntry TypeScheme]]
forall k a. Map k a -> [a]
Map.elems Map String [VarEntry TypeScheme]
env
  where
    freeVarsInScheme :: VarEntry TypeScheme -> Set TyVar
freeVarsInScheme VarEntry TypeScheme
entry = 
      let Forall [TyVar]
vs [Constraint]
_ Type
t = VarEntry TypeScheme -> TypeScheme
forall a. VarEntry a -> a
veValue VarEntry TypeScheme
entry
      in Type -> Set TyVar
freeTyVars Type
t Set TyVar -> Set TyVar -> Set TyVar
forall a. Ord a => Set a -> Set a -> Set a
`Set.difference` [TyVar] -> Set TyVar
forall a. Ord a => [a] -> Set a
Set.fromList [TyVar]
vs

-- | Generalize a type to a type scheme (without constraints)
-- Generalize all free type variables that are not in the environment
generalize :: TypeEnv -> Type -> TypeScheme
generalize :: TypeEnv -> Type -> TypeScheme
generalize TypeEnv
env Type
t =
  let envFreeVars :: Set TyVar
envFreeVars = TypeEnv -> Set TyVar
freeVarsInEnv TypeEnv
env
      typeFreeVars :: Set TyVar
typeFreeVars = Type -> Set TyVar
freeTyVars Type
t
      genVars :: [TyVar]
genVars = Set TyVar -> [TyVar]
forall a. Set a -> [a]
Set.toList (Set TyVar -> [TyVar]) -> Set TyVar -> [TyVar]
forall a b. (a -> b) -> a -> b
$ Set TyVar
typeFreeVars Set TyVar -> Set TyVar -> Set TyVar
forall a. Ord a => Set a -> Set a -> Set a
`Set.difference` Set TyVar
envFreeVars
  in [TyVar] -> [Constraint] -> Type -> TypeScheme
Forall [TyVar]
genVars [] Type
t

-- | Instantiate a type scheme with fresh type variables
-- Returns a tuple of (constraints, instantiated type, fresh variable counter)
instantiate :: TypeScheme -> Int -> ([Constraint], Type, Int)
instantiate :: TypeScheme -> Int -> ([Constraint], Type, Int)
instantiate (Forall [TyVar]
vs [Constraint]
cs Type
t) Int
counter =
  let freshVars :: [(TyVar, Type)]
freshVars = (TyVar -> Int -> (TyVar, Type))
-> [TyVar] -> [Int] -> [(TyVar, Type)]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\TyVar
v Int
i -> (TyVar
v, TyVar -> Type
TVar (String -> Int -> TyVar
freshTyVar String
"t" (Int
counter Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
i)))) [TyVar]
vs [Int
0..]
      substType :: Type
substType = ((TyVar, Type) -> Type -> Type) -> Type -> [(TyVar, 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 (\(TyVar
old, Type
new) Type
acc -> TyVar -> Type -> Type -> Type
substVar TyVar
old Type
new Type
acc) Type
t [(TyVar, Type)]
freshVars
      substCs :: [Constraint]
substCs = (Constraint -> Constraint) -> [Constraint] -> [Constraint]
forall a b. (a -> b) -> [a] -> [b]
map ([(TyVar, Type)] -> Constraint -> Constraint
substConstraint [(TyVar, Type)]
freshVars) [Constraint]
cs
  in ([Constraint]
substCs, Type
substType, Int
counter Int -> Int -> Int
forall a. Num a => a -> a -> a
+ [TyVar] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [TyVar]
vs)
  where
    substConstraint :: [(TyVar, Type)] -> Constraint -> Constraint
    substConstraint :: [(TyVar, Type)] -> Constraint -> Constraint
substConstraint [(TyVar, Type)]
vars (Constraint String
cls Type
ty) =
      String -> Type -> Constraint
Constraint String
cls (((TyVar, Type) -> Type -> Type) -> Type -> [(TyVar, 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 (\(TyVar
old, Type
new) Type
acc -> TyVar -> Type -> Type -> Type
substVar TyVar
old Type
new Type
acc) Type
ty [(TyVar, Type)]
vars)
    substVar :: TyVar -> Type -> Type -> Type
    substVar :: TyVar -> Type -> Type -> Type
substVar TyVar
_ Type
_ Type
TInt = Type
TInt
    substVar TyVar
_ Type
_ Type
TMathExpr = Type
TMathExpr
    substVar TyVar
_ Type
_ Type
TPolyExpr = Type
TPolyExpr
    substVar TyVar
_ Type
_ Type
TTermExpr = Type
TTermExpr
    substVar TyVar
_ Type
_ Type
TSymbolExpr = Type
TSymbolExpr
    substVar TyVar
_ Type
_ Type
TIndexExpr = Type
TIndexExpr
    substVar TyVar
_ Type
_ Type
TFloat = Type
TFloat
    substVar TyVar
_ Type
_ Type
TBool = Type
TBool
    substVar TyVar
_ Type
_ Type
TChar = Type
TChar
    substVar TyVar
_ Type
_ Type
TString = Type
TString
    substVar TyVar
old Type
new (TVar TyVar
v)
      | TyVar
v TyVar -> TyVar -> Bool
forall a. Eq a => a -> a -> Bool
== TyVar
old = Type
new
      | Bool
otherwise = TyVar -> Type
TVar TyVar
v
    substVar TyVar
old Type
new (TTuple [Type]
ts) = [Type] -> Type
TTuple ((Type -> Type) -> [Type] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map (TyVar -> Type -> Type -> Type
substVar TyVar
old Type
new) [Type]
ts)
    substVar TyVar
old Type
new (TCollection Type
t') = Type -> Type
TCollection (TyVar -> Type -> Type -> Type
substVar TyVar
old Type
new Type
t')
    substVar TyVar
old Type
new (TInductive String
name [Type]
ts) = String -> [Type] -> Type
TInductive String
name ((Type -> Type) -> [Type] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map (TyVar -> Type -> Type -> Type
substVar TyVar
old Type
new) [Type]
ts)
    substVar TyVar
old Type
new (TTensor Type
t') = Type -> Type
TTensor (TyVar -> Type -> Type -> Type
substVar TyVar
old Type
new Type
t')
    substVar TyVar
old Type
new (THash Type
k Type
v) = Type -> Type -> Type
THash (TyVar -> Type -> Type -> Type
substVar TyVar
old Type
new Type
k) (TyVar -> Type -> Type -> Type
substVar TyVar
old Type
new Type
v)
    substVar TyVar
old Type
new (TMatcher Type
t') = Type -> Type
TMatcher (TyVar -> Type -> Type -> Type
substVar TyVar
old Type
new Type
t')
    substVar TyVar
old Type
new (TFun Type
t1 Type
t2) = Type -> Type -> Type
TFun (TyVar -> Type -> Type -> Type
substVar TyVar
old Type
new Type
t1) (TyVar -> Type -> Type -> Type
substVar TyVar
old Type
new Type
t2)
    substVar TyVar
old Type
new (TIO Type
t') = Type -> Type
TIO (TyVar -> Type -> Type -> Type
substVar TyVar
old Type
new Type
t')
    substVar TyVar
old Type
new (TIORef Type
t') = Type -> Type
TIORef (TyVar -> Type -> Type -> Type
substVar TyVar
old Type
new Type
t')
    substVar TyVar
_ Type
_ Type
TPort = Type
TPort
    substVar TyVar
_ Type
_ Type
TAny = Type
TAny

--------------------------------------------------------------------------------
-- Class Environment
--------------------------------------------------------------------------------

-- | Class environment: maps class names to class info and instances
data ClassEnv = ClassEnv
  { ClassEnv -> Map String ClassInfo
classEnvClasses   :: Map String ClassInfo      -- ^ Class definitions
  , ClassEnv -> Map String [InstanceInfo]
classEnvInstances :: Map String [InstanceInfo] -- ^ Instances per class
  } deriving (ClassEnv -> ClassEnv -> Bool
(ClassEnv -> ClassEnv -> Bool)
-> (ClassEnv -> ClassEnv -> Bool) -> Eq ClassEnv
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: ClassEnv -> ClassEnv -> Bool
== :: ClassEnv -> ClassEnv -> Bool
$c/= :: ClassEnv -> ClassEnv -> Bool
/= :: ClassEnv -> ClassEnv -> Bool
Eq, Int -> ClassEnv -> ShowS
[ClassEnv] -> ShowS
ClassEnv -> String
(Int -> ClassEnv -> ShowS)
-> (ClassEnv -> String) -> ([ClassEnv] -> ShowS) -> Show ClassEnv
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> ClassEnv -> ShowS
showsPrec :: Int -> ClassEnv -> ShowS
$cshow :: ClassEnv -> String
show :: ClassEnv -> String
$cshowList :: [ClassEnv] -> ShowS
showList :: [ClassEnv] -> ShowS
Show)

-- | Empty class environment
emptyClassEnv :: ClassEnv
emptyClassEnv :: ClassEnv
emptyClassEnv = Map String ClassInfo -> Map String [InstanceInfo] -> ClassEnv
ClassEnv Map String ClassInfo
forall k a. Map k a
Map.empty Map String [InstanceInfo]
forall k a. Map k a
Map.empty

-- | Add a class to the environment
addClass :: String -> ClassInfo -> ClassEnv -> ClassEnv
addClass :: String -> ClassInfo -> ClassEnv -> ClassEnv
addClass String
name ClassInfo
info (ClassEnv Map String ClassInfo
classes Map String [InstanceInfo]
insts) =
  Map String ClassInfo -> Map String [InstanceInfo] -> ClassEnv
ClassEnv (String -> ClassInfo -> Map String ClassInfo -> Map String ClassInfo
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert String
name ClassInfo
info Map String ClassInfo
classes) Map String [InstanceInfo]
insts

-- | Add an instance to the environment
addInstance :: String -> InstanceInfo -> ClassEnv -> ClassEnv
addInstance :: String -> InstanceInfo -> ClassEnv -> ClassEnv
addInstance String
className InstanceInfo
inst (ClassEnv Map String ClassInfo
classes Map String [InstanceInfo]
insts) =
  Map String ClassInfo -> Map String [InstanceInfo] -> ClassEnv
ClassEnv Map String ClassInfo
classes (([InstanceInfo] -> [InstanceInfo] -> [InstanceInfo])
-> String
-> [InstanceInfo]
-> Map String [InstanceInfo]
-> Map String [InstanceInfo]
forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
Map.insertWith [InstanceInfo] -> [InstanceInfo] -> [InstanceInfo]
forall a. [a] -> [a] -> [a]
(++) String
className [InstanceInfo
inst] Map String [InstanceInfo]
insts)

-- | Look up a class definition
lookupClass :: String -> ClassEnv -> Maybe ClassInfo
lookupClass :: String -> ClassEnv -> Maybe ClassInfo
lookupClass String
name (ClassEnv Map String ClassInfo
classes Map String [InstanceInfo]
_) = String -> Map String ClassInfo -> Maybe ClassInfo
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup String
name Map String ClassInfo
classes

-- | Look up instances for a class
lookupInstances :: String -> ClassEnv -> [InstanceInfo]
lookupInstances :: String -> ClassEnv -> [InstanceInfo]
lookupInstances String
name (ClassEnv Map String ClassInfo
_ Map String [InstanceInfo]
insts) = [InstanceInfo]
-> String -> Map String [InstanceInfo] -> [InstanceInfo]
forall k a. Ord k => a -> k -> Map k a -> a
Map.findWithDefault [] String
name Map String [InstanceInfo]
insts

-- | Convert class environment to list
classEnvToList :: ClassEnv -> [(String, ClassInfo)]
classEnvToList :: ClassEnv -> [(String, ClassInfo)]
classEnvToList (ClassEnv Map String ClassInfo
classes Map String [InstanceInfo]
_) = Map String ClassInfo -> [(String, ClassInfo)]
forall k a. Map k a -> [(k, a)]
Map.toList Map String ClassInfo
classes

-- | Merge two class environments
-- The second environment's definitions take precedence in case of conflicts
mergeClassEnv :: ClassEnv -> ClassEnv -> ClassEnv
mergeClassEnv :: ClassEnv -> ClassEnv -> ClassEnv
mergeClassEnv (ClassEnv Map String ClassInfo
classes1 Map String [InstanceInfo]
insts1) (ClassEnv Map String ClassInfo
classes2 Map String [InstanceInfo]
insts2) =
  Map String ClassInfo -> Map String [InstanceInfo] -> ClassEnv
ClassEnv
    (Map String ClassInfo
-> Map String ClassInfo -> Map String ClassInfo
forall k a. Ord k => Map k a -> Map k a -> Map k a
Map.union Map String ClassInfo
classes2 Map String ClassInfo
classes1)  -- classes2 takes precedence
    (([InstanceInfo] -> [InstanceInfo] -> [InstanceInfo])
-> Map String [InstanceInfo]
-> Map String [InstanceInfo]
-> Map String [InstanceInfo]
forall k a. Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a
Map.unionWith [InstanceInfo] -> [InstanceInfo] -> [InstanceInfo]
forall a. [a] -> [a] -> [a]
(++) Map String [InstanceInfo]
insts2 Map String [InstanceInfo]
insts1)  -- Combine instance lists

--------------------------------------------------------------------------------
-- Pattern Type Environment
--------------------------------------------------------------------------------

-- | Empty pattern type environment
emptyPatternEnv :: PatternTypeEnv
emptyPatternEnv :: PatternTypeEnv
emptyPatternEnv = Map String TypeScheme -> PatternTypeEnv
PatternTypeEnv Map String TypeScheme
forall k a. Map k a
Map.empty

-- | Extend the pattern type environment with a new binding
extendPatternEnv :: String -> TypeScheme -> PatternTypeEnv -> PatternTypeEnv
extendPatternEnv :: String -> TypeScheme -> PatternTypeEnv -> PatternTypeEnv
extendPatternEnv String
name TypeScheme
scheme (PatternTypeEnv Map String TypeScheme
env) = Map String TypeScheme -> PatternTypeEnv
PatternTypeEnv (Map String TypeScheme -> PatternTypeEnv)
-> Map String TypeScheme -> PatternTypeEnv
forall a b. (a -> b) -> a -> b
$ String
-> TypeScheme -> Map String TypeScheme -> Map String TypeScheme
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert String
name TypeScheme
scheme Map String TypeScheme
env

-- | Look up a pattern constructor/function in the environment
lookupPatternEnv :: String -> PatternTypeEnv -> Maybe TypeScheme
lookupPatternEnv :: String -> PatternTypeEnv -> Maybe TypeScheme
lookupPatternEnv String
name (PatternTypeEnv Map String TypeScheme
env) = String -> Map String TypeScheme -> Maybe TypeScheme
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup String
name Map String TypeScheme
env

-- | Convert pattern type environment to list
patternEnvToList :: PatternTypeEnv -> [(String, TypeScheme)]
patternEnvToList :: PatternTypeEnv -> [(String, TypeScheme)]
patternEnvToList (PatternTypeEnv Map String TypeScheme
env) = Map String TypeScheme -> [(String, TypeScheme)]
forall k a. Map k a -> [(k, a)]
Map.toList Map String TypeScheme
env