import LeanPython.AST.Types

import Std.Data.HashMap

set_option autoImplicit false

namespace LeanPython.Runtime

open LeanPython.AST (Arguments Stmt Expr)

open LeanPython.Lexer (SourceSpan)

abbrev HeapRef := Nat

inductive Value where

| none

| bool : Bool → Value

| int : Int → Value

| float : Float → Value

| str : String → Value

| bytes : ByteArray → Value

| list : HeapRef → Value

| tuple : Array Value → Value

| dict : HeapRef → Value

| set : HeapRef → Value

| function : HeapRef → Value

| builtin : String → Value

| ellipsis

| boundMethod : Value → String → Value

| exception : String → String → Value

| generator : HeapRef → Value

| coroutine : Value → Value

| classObj : HeapRef → Value

| instance : HeapRef → Value

| superObj : Value → Value → Value

| staticMethod : Value → Value

| classMethod : Value → Value

| property : Value → Option Value → Option Value → Value

| module : HeapRef → Value

instance : Inhabited Value where

default := .none

abbrev Scope := Std.HashMap String Value

abbrev ClosureEnv := Array Scope

structure FuncData where

name : String

params : Arguments

body : List Stmt

defaults : Array Value

kwDefaults : Array (Option Value)

closure : ClosureEnv

isGenerator : Bool

definingModule : Option String := none

isAsync : Bool := false

structure ClassData where

name : String

bases : Array Value

mro : Array Value

ns : Std.HashMap String Value

slots : Option (Array String) := none

structure InstanceData where

cls : Value

attrs : Std.HashMap String Value

wrappedValue : Option Value := none

structure ModuleData where

name : String -- __name__

file : Option String -- __file__

package : Option String -- __package__

ns : Std.HashMap String Value -- module namespace

allNames : Option (Array String) -- __all__ if defined

inductive HeapObject where

| listObj : Array Value → HeapObject

| dictObj : Array (Value × Value) → HeapObject

| setObj : Array Value → HeapObject

| funcObj : FuncData → HeapObject

| generatorObj : Array Value → Nat → HeapObject

| classObjData : ClassData → HeapObject

| instanceObjData : InstanceData → HeapObject

| moduleObj : ModuleData → HeapObject

inductive RuntimeError where

| nameError : String → RuntimeError

| typeError : String → RuntimeError

| valueError : String → RuntimeError

| indexError : String → RuntimeError

| keyError : String → RuntimeError

| zeroDivision : String → RuntimeError

| assertionError : String → RuntimeError

| attributeError : String → RuntimeError

| overflowError : String → RuntimeError

| stopIteration : RuntimeError

| notImplemented : String → RuntimeError

| runtimeError : String → RuntimeError

| importError : String → RuntimeError

| moduleNotFound : String → RuntimeError

| customError : String → String → List String → RuntimeError -- typeName, message, parentTypeNames

deriving Repr

instance : ToString RuntimeError where

toString

| .nameError s => s!"NameError: {s}"

| .typeError s => s!"TypeError: {s}"

| .valueError s => s!"ValueError: {s}"

| .indexError s => s!"IndexError: {s}"

| .keyError s => s!"KeyError: {s}"

| .zeroDivision s => s!"ZeroDivisionError: {s}"

| .assertionError s => s!"AssertionError: {s}"

| .attributeError s => s!"AttributeError: {s}"

| .overflowError s => s!"OverflowError: {s}"

| .stopIteration => "StopIteration"

| .notImplemented s => s!"NotImplementedError: {s}"

| .runtimeError s => s!"RuntimeError: {s}"

| .importError s => s!"ImportError: {s}"

| .moduleNotFound s => s!"ModuleNotFoundError: No module named '{s}'"

| .customError tn msg _ => if msg.isEmpty then tn else s!"{tn}: {msg}"

partial def Value.beq : Value → Value → Bool

| .none, .none => true

| .bool a, .bool b => a == b

| .int a, .int b => a == b

| .float a, .float b => a == b

| .str a, .str b => a == b

| .bytes a, .bytes b => a == b

| .list a, .list b => a == b

| .tuple a, .tuple b => a.size == b.size && (List.range a.size).all fun i =>

match a[i]?, b[i]? with

| some va, some vb => Value.beq va vb

| _, _ => false

| .dict a, .dict b => a == b

| .set a, .set b => a == b

| .function a, .function b => a == b

| .builtin a, .builtin b => a == b

| .ellipsis, .ellipsis => true

| .boundMethod _ _, _ => false

| _, .boundMethod _ _ => false

| .exception a1 a2, .exception b1 b2 => a1 == b1 && a2 == b2

| .generator a, .generator b => a == b

| .coroutine a, .coroutine b => Value.beq a b

| .classObj a, .classObj b => a == b

| .instance a, .instance b => a == b

| .superObj _ _, _ => false

| _, .superObj _ _ => false

| .staticMethod a, .staticMethod b => Value.beq a b

| .classMethod a, .classMethod b => Value.beq a b

| .property a1 _ _, .property b1 _ _ => Value.beq a1 b1

| .staticMethod _, _ => false

| _, .staticMethod _ => false

| .classMethod _, _ => false

| _, .classMethod _ => false

| .property _ _ _, _ => false

| _, .property _ _ _ => false

| .module a, .module b => a == b

| .module _, _ => false

| _, .module _ => false

-- Cross-type: bool/int interop (Python: True == 1, False == 0)

| .bool a, .int b => (if a then 1 else 0) == b

| .int a, .bool b => a == (if b then 1 else 0)

-- int/float comparison

| .int a, .float b => Float.ofInt a == b

| .float a, .int b => a == Float.ofInt b

| .bool a, .float b => Float.ofInt (if a then 1 else 0) == b

| .float a, .bool b => a == Float.ofInt (if b then 1 else 0)

| _, _ => false

instance : BEq Value where

beq := Value.beq

partial def Value.toStr : Value → String

| .none => "None"

| .bool true => "True"

| .bool false => "False"

| .int n => toString n

| .float f =>

let s := toString f

s

| .str s => s

-- str(exception) returns just the message (not the type name)

| .exception _typeName msg => msg

| .bytes _b => "b'...'"

| .list _ => "[...]"

| .tuple elems =>

if elems.size == 1 then

s!"({Value.toStr elems[0]!},)"

else

"(" ++ ", ".intercalate (elems.toList.map Value.toStr) ++ ")"

| .dict _ => "{...}"

| .set _ => "{...}"

| .function _ => "<function>"

| .builtin name => s!"<built-in function {name}>"

| .ellipsis => "Ellipsis"

| .boundMethod _ method => s!"<bound method {method}>"

| .generator _ => "<generator object>"

| .coroutine _ => "<coroutine object>"

| .classObj _ => "<class>"

| .instance _ => "<instance>"

| .superObj _ _ => "<super>"

| .staticMethod _ => "<staticmethod object>"

| .classMethod _ => "<classmethod object>"

| .property _ _ _ => "<property object>"

| .module _ => "<module>"

partial def Value.toRepr : Value → String

| .none => "None"

| .bool true => "True"

| .bool false => "False"

| .int n => toString n

| .float f => toString f

| .str s => s!"'{s}'"

| .bytes _b => "b'...'"

| .list _ => "[...]"

| .tuple elems =>

if elems.size == 1 then

s!"({Value.toRepr elems[0]!},)"

else

"(" ++ ", ".intercalate (elems.toList.map Value.toRepr) ++ ")"

| .dict _ => "{...}"

| .set _ => "{...}"

| .function _ => "<function>"

| .builtin name => s!"<built-in function {name}>"

| .ellipsis => "Ellipsis"

| .boundMethod _ method => s!"<bound method {method}>"

| .exception typeName msg => if msg.isEmpty then typeName else s!"{typeName}('{msg}')"

| .generator _ => "<generator object>"

| .coroutine _ => "<coroutine object>"

| .classObj _ => "<class>"

| .instance _ => "<instance>"

| .superObj _ _ => "<super>"

| .staticMethod _ => "<staticmethod object>"

| .classMethod _ => "<classmethod object>"

| .property _ _ _ => "<property object>"

| .module _ => "<module>"

instance : ToString Value where

toString := Value.toStr

def runtimeErrorTypeName : RuntimeError → String

| .nameError _ => "NameError"

| .typeError _ => "TypeError"

| .valueError _ => "ValueError"

| .indexError _ => "IndexError"

| .keyError _ => "KeyError"

| .zeroDivision _ => "ZeroDivisionError"

| .assertionError _ => "AssertionError"

| .attributeError _ => "AttributeError"

| .overflowError _ => "OverflowError"

| .stopIteration => "StopIteration"

| .notImplemented _ => "NotImplementedError"

| .runtimeError _ => "RuntimeError"

| .importError _ => "ImportError"

| .moduleNotFound _ => "ModuleNotFoundError"

| .customError tn _ _ => tn

def runtimeErrorMessage : RuntimeError → String

| .nameError s | .typeError s | .valueError s | .indexError s

| .keyError s | .zeroDivision s | .assertionError s

| .attributeError s | .overflowError s | .notImplemented s

| .runtimeError s | .importError s | .moduleNotFound s => s

| .stopIteration => ""

| .customError _ msg _ => msg

def exceptionMatches (errorTypeName handlerTypeName : String) : Bool :=

if handlerTypeName == "BaseException" then true

else if handlerTypeName == "Exception" then

-- Exception catches everything except BaseException-only subtypes

errorTypeName != "SystemExit" && errorTypeName != "KeyboardInterrupt" &&

errorTypeName != "GeneratorExit"

else

errorTypeName == handlerTypeName

def customExceptionMatches (error : RuntimeError) (handlerTypeName : String) : Bool :=

match error with

| .customError tn _ parents =>

if handlerTypeName == "BaseException" || handlerTypeName == "Exception" then true

else tn == handlerTypeName || parents.any (· == handlerTypeName)

| other => exceptionMatches (runtimeErrorTypeName other) handlerTypeName

def builtinNames : List String :=

["print", "len", "range", "type", "int", "str", "bool", "float",

"list", "tuple", "set", "dict", "isinstance", "repr", "abs",

"min", "max", "sorted", "reversed", "enumerate", "zip",

"sum", "any", "all", "hash", "id", "input", "ord", "chr",

"hex", "oct", "bin", "round", "pow", "divmod", "map", "filter",

"iter", "next", "hasattr", "getattr", "setattr", "callable",

"issubclass", "super", "object", "bytes", "bytearray",

"memoryview", "frozenset", "complex", "slice",

"staticmethod", "classmethod", "property",

-- Dataclass

"dataclass",

-- Exception classes

"ValueError", "TypeError", "KeyError", "IndexError",

"RuntimeError", "ZeroDivisionError", "AssertionError",

"AttributeError", "OverflowError", "StopIteration",

"NotImplementedError", "Exception", "BaseException",

"NameError", "OSError", "IOError", "FileNotFoundError",

"ImportError", "ModuleNotFoundError", "SystemExit",

"KeyboardInterrupt", "GeneratorExit",

"CancelledError", "TimeoutError"]

def isBuiltinName (name : String) : Bool :=

builtinNames.contains name

def isBuiltinExceptionName (name : String) : Bool :=

name == "Exception" || name == "BaseException" ||

name == "ValueError" || name == "TypeError" ||

name == "KeyError" || name == "IndexError" ||

name == "RuntimeError" || name == "ZeroDivisionError" ||

name == "AssertionError" || name == "AttributeError" ||

name == "OverflowError" || name == "StopIteration" ||

name == "NotImplementedError" || name == "NameError" ||

name == "OSError" || name == "IOError" ||

name == "FileNotFoundError" || name == "ImportError" ||

name == "ModuleNotFoundError" || name == "SystemExit" ||

name == "KeyboardInterrupt" || name == "GeneratorExit" ||

name == "CancelledError" || name == "TimeoutError"

end LeanPython.Runtime