From c82b41a4dfbadd693ff2f7e7b00d14aac62be585 Mon Sep 17 00:00:00 2001
From: LordBaryhobal <lordbaryhobal@gmail.com>
Date: Thu, 28 May 2026 17:29:37 +0200
Subject: [PATCH] feat(checker): add environment manager

adapted from Pebble
---
 midas/checker/environment.py | 52 ++++++++++++++++++++++++++++++++++++
 1 file changed, 52 insertions(+)
 create mode 100644 midas/checker/environment.py

diff --git a/midas/checker/environment.py b/midas/checker/environment.py
new file mode 100644
index 0000000..dc5cac8
--- /dev/null
+++ b/midas/checker/environment.py
@@ -0,0 +1,52 @@
+from __future__ import annotations
+
+from typing import Optional
+
+from midas.checker.types import Type
+
+
+class Environment:
+    def __init__(self, enclosing: Optional[Environment] = None) -> None:
+        self.enclosing: Optional[Environment] = enclosing
+        self.values: dict[str, Type] = {}
+
+    def define(self, name: str, value: Type):
+        self.values[name] = value
+
+    def get(self, name: str) -> Optional[Type]:
+        if name in self.values:
+            return self.values[name]
+        if self.enclosing is not None:
+            return self.enclosing.get(name)
+        # raise NameError(f"Undefined variable '{name}'")
+        return None
+
+    def assign(self, name: str, value: Type) -> bool:
+        if name not in self.values:
+            if self.enclosing is None:
+                return False
+            if self.enclosing.assign(name, value):
+                return True
+        self.values[name] = value
+        return True
+
+    def clear(self):
+        self.values = {}
+
+    def get_at(self, distance: int, name: str) -> Optional[Type]:
+        return self.ancestor(distance).values.get(name)
+
+    def assign_at(self, distance: int, name: str, value: Type):
+        self.ancestor(distance).values[name] = value
+
+    def ancestor(self, distance: int) -> Environment:
+        env: Environment = self
+        for _ in range(distance):
+            assert env.enclosing is not None
+            env = env.enclosing
+        return env
+
+    def flat_dict(self) -> dict:
+        if self.enclosing is None:
+            return self.values
+        return self.enclosing.flat_dict() | self.values