chore(compiler): do not store symbol environments on AST (#3712)

## Checklist

- [ ] Title matches [Winglang's style guide](https://www.winglang.io/contributing/start-here/pull_requests#how-are-pull-request-titles-formatted)
- [ ] Description explains motivation and solution
- [ ] Tests added (always)
- [ ] Docs updated (only required for features)
- [ ] Added `pr/e2e-full` label if this feature requires end-to-end testing

*By submitting this pull request, I confirm that my contribution is made under the terms of the [Wing Cloud Contribution License](https://github.com/winglang/wing/blob/main/CONTRIBUTION_LICENSE.md)*.

libs/wingc/src/ast.rs

@@ -1,17 +1,15 @@
-use std::cell::RefCell;
 use std::fmt::{Debug, Display};
 use std::hash::{Hash, Hasher};
 use std::sync::atomic::{AtomicUsize, Ordering};
 
-use derivative::Derivative;
 use indexmap::{Equivalent, IndexMap, IndexSet};
 use itertools::Itertools;
 
 use crate::diagnostic::WingSpan;
-use crate::type_check::symbol_env::SymbolEnvRef;
 use crate::type_check::CLOSURE_CLASS_HANDLE_METHOD;
 
 static EXPR_COUNTER: AtomicUsize = AtomicUsize::new(0);
+static SCOPE_COUNTER: AtomicUsize = AtomicUsize::new(0);
 
 #[derive(Debug, Eq, Clone)]
 pub struct Symbol {
@@ -584,6 +582,7 @@ impl Spanned for CalleeKind {
 	}
 }
 
+// do not derive Default, we want to be explicit about generating ids
 #[derive(Debug)]
 pub struct Expr {
 	/// An identifier that is unique among all expressions in the AST.
@@ -597,7 +596,6 @@ pub struct Expr {
 impl Expr {
 	pub fn new(kind: ExprKind, span: WingSpan) -> Self {
 		let id = EXPR_COUNTER.fetch_add(1, Ordering::SeqCst);
-
 		Self { id, kind, span }
 	}
 
@@ -655,33 +653,29 @@ pub enum InterpolatedStringPart {
 	Expr(Expr),
 }
 
-#[derive(Derivative, Default)]
-#[derivative(Debug)]
+pub type ScopeId = usize;
+
+// do not derive Default, as we want to explicitly generate IDs
+#[derive(Debug)]
 pub struct Scope {
+	/// An identifier that is unique among all scopes in the AST.
+	pub id: ScopeId,
 	pub statements: Vec<Stmt>,
 	pub span: WingSpan,
-	#[derivative(Debug = "ignore")]
-	pub env: RefCell<Option<SymbolEnvRef>>, // None after parsing, set to Some during type checking phase
 }
 
 impl Scope {
-	pub fn new(statements: Vec<Stmt>, span: WingSpan) -> Self {
+	pub fn empty() -> Self {
 		Self {
-			statements,
-			span,
-			env: RefCell::new(None),
+			id: SCOPE_COUNTER.fetch_add(1, Ordering::SeqCst),
+			statements: vec![],
+			span: WingSpan::default(),
 		}
 	}
 
-	pub fn set_env(&self, new_env: SymbolEnvRef) {
-		let mut env = self.env.borrow_mut();
-		assert!((*env).is_none());
-		*env = Some(new_env);
-	}
-
-	pub fn reset_env(&self) {
-		let mut env = self.env.borrow_mut();
-		*env = None;
+	pub fn new(statements: Vec<Stmt>, span: WingSpan) -> Self {
+		let id = SCOPE_COUNTER.fetch_add(1, Ordering::SeqCst);
+		Self { id, statements, span }
 	}
 }
 

libs/wingc/src/closure_transform.rs

@@ -92,9 +92,9 @@ impl Fold for ClosureTransformer {
 		}
 
 		Scope {
+			id: node.id,
 			statements,
 			span: node.span,
-			env: node.env,
 		}
 	}
 

libs/wingc/src/fold.rs

@@ -70,9 +70,9 @@ where
 	F: Fold + ?Sized,
 {
 	Scope {
+		id: node.id,
 		statements: node.statements.into_iter().map(|stmt| f.fold_stmt(stmt)).collect(),
 		span: node.span,
-		env: node.env,
 	}
 }
 

libs/wingc/src/jsify.rs

@@ -105,17 +105,16 @@ impl<'a> JSifier<'a> {
 			visit_ctx: &mut visit_ctx,
 			lifts: None,
 		};
-		jsify_context
-			.visit_ctx
-			.push_env(scope.env.borrow().as_ref().unwrap().get_ref());
+		jsify_context.visit_ctx.push_env(self.types.get_scope_env(&scope));
 		for statement in scope.statements.iter().sorted_by(|a, b| match (&a.kind, &b.kind) {
 			// Put type definitions first so JS won't complain of unknown types
 			(StmtKind::Class(AstClass { .. }), StmtKind::Class(AstClass { .. })) => Ordering::Equal,
 			(StmtKind::Class(AstClass { .. }), _) => Ordering::Less,
 			(_, StmtKind::Class(AstClass { .. })) => Ordering::Greater,
 			_ => Ordering::Equal,
 		}) {
-			let s = self.jsify_statement(scope.env.borrow().as_ref().unwrap(), statement, &mut jsify_context); // top level statements are always preflight
+			let scope_env = self.types.get_scope_env(&scope);
+			let s = self.jsify_statement(&scope_env, statement, &mut jsify_context); // top level statements are always preflight
 			if let StmtKind::Bring {
 				identifier: _,
 				source: _,
@@ -212,9 +211,10 @@ impl<'a> JSifier<'a> {
 		CompilationContext::set(CompilationPhase::Jsifying, &scope.span);
 		let mut code = CodeMaker::default();
 
-		ctx.visit_ctx.push_env(scope.env.borrow().as_ref().unwrap().get_ref());
+		let scope_env = self.types.get_scope_env(&scope);
+		ctx.visit_ctx.push_env(scope_env);
 		for statement in scope.statements.iter() {
-			let statement_code = self.jsify_statement(scope.env.borrow().as_ref().unwrap(), statement, ctx);
+			let statement_code = self.jsify_statement(&scope_env, statement, ctx);
 			code.add_code(statement_code);
 		}
 		ctx.visit_ctx.pop_env();

libs/wingc/src/lib.rs

@@ -53,7 +53,6 @@ pub mod jsify;
 mod lifting;
 pub mod lsp;
 pub mod parser;
-mod reset;
 pub mod type_check;
 mod type_check_assert;
 pub mod visit;
@@ -167,9 +166,8 @@ pub fn type_check(
 	jsii_types: &mut TypeSystem,
 	jsii_imports: &mut Vec<JsiiImportSpec>,
 ) {
-	assert!(scope.env.borrow().is_none(), "Scope should not have an env yet");
 	let env = types.add_symbol_env(SymbolEnv::new(None, types.void(), false, false, Phase::Preflight, 0));
-	scope.set_env(env);
+	types.set_scope_env(scope, env);
 
 	// note: Globals are emitted here and wrapped in "{ ... }" blocks. Wrapping makes these emissions, actual
 	// statements and not expressions. this makes the runtime panic if these are used in place of expressions.
@@ -248,20 +246,24 @@ pub fn type_check(
 		types,
 	);
 
+	let mut scope_env = types.get_scope_env(&scope);
 	let mut tc = TypeChecker::new(types, file_path, jsii_types, jsii_imports);
-	tc.add_globals(scope);
+	tc.add_module_to_env(
+		&mut scope_env,
+		WINGSDK_ASSEMBLY_NAME.to_string(),
+		vec![WINGSDK_STD_MODULE.to_string()],
+		&Symbol::global(WINGSDK_STD_MODULE),
+		None,
+	);
 
 	tc.type_check_file(file_path, scope);
 }
 
 // TODO: refactor this (why is scope needed?) (move to separate module?)
 fn add_builtin(name: &str, typ: Type, scope: &mut Scope, types: &mut Types) {
 	let sym = Symbol::global(name);
-	scope
-		.env
-		.borrow_mut()
-		.as_mut()
-		.unwrap()
+	let mut scope_env = types.get_scope_env(&scope);
+	scope_env
 		.define(
 			&sym,
 			SymbolKind::make_free_variable(sym.clone(), types.add_type(typ), false, Phase::Independent),

libs/wingc/src/lifting.rs

@@ -243,7 +243,7 @@ impl<'a> Visit<'a> for LiftVisitor<'a> {
 				self.ctx.push_function_definition(
 					&node.name,
 					&node.signature.phase,
-					scope.env.borrow().as_ref().unwrap().get_ref(),
+					self.jsify.types.get_scope_env(&scope),
 				);
 
 				visit::visit_function_definition(self, node);
@@ -280,7 +280,7 @@ impl<'a> Visit<'a> for LiftVisitor<'a> {
 		// because this is the environment in which we want to resolve references
 		// as oppose to the environment of the class definition itself.
 		let init_env = if let FunctionBody::Statements(ref s) = node.initializer.body {
-			Some(s.env.borrow().as_ref().unwrap().get_ref())
+			Some(self.jsify.types.get_scope_env(&s))
 		} else {
 			None
 		};
@@ -326,7 +326,7 @@ impl<'a> Visit<'a> for LiftVisitor<'a> {
 	}
 
 	fn visit_scope(&mut self, node: &'a Scope) {
-		self.ctx.push_env(node.env.borrow().as_ref().unwrap().get_ref());
+		self.ctx.push_env(self.jsify.types.get_scope_env(&node));
 		visit::visit_scope(self, node);
 		self.ctx.pop_env();
 	}

libs/wingc/src/lsp/completions.rs

@@ -46,8 +46,7 @@ pub fn on_completion(params: lsp_types::CompletionParams) -> CompletionResponse
 			let file_id = file.to_str().expect("File path must be valid utf8");
 			let root_ts_node = project_data.trees.get(&file).expect("tree not found").root_node();
 			let root_scope = project_data.asts.get(&file).expect("ast not found");
-			let root_env = root_scope.env.borrow();
-			let root_env = root_env.as_ref().expect("The root scope must have an environment");
+			let root_env = types.get_scope_env(root_scope);
 			let contents = project_data.files.get_file(&file).expect("file not found");
 
 			// get all character from file_data.contents up to the current position
@@ -102,8 +101,7 @@ pub fn on_completion(params: lsp_types::CompletionParams) -> CompletionResponse
 			);
 			scope_visitor.visit();
 
-			let found_env = scope_visitor.found_scope.env.borrow();
-			let found_env = found_env.as_ref().unwrap();
+			let found_env = types.get_scope_env(&scope_visitor.found_scope);
 
 			// references have a complicated hierarchy, so it's useful to know the nearest non-reference parent
 			let mut nearest_non_reference_parent = node_to_complete.parent();
@@ -196,7 +194,7 @@ pub fn on_completion(params: lsp_types::CompletionParams) -> CompletionResponse
 
 				if let Some(nearest_type_annotation) = scope_visitor.nearest_type_annotation {
 					if let TypeAnnotationKind::UserDefined(udt) = &nearest_type_annotation.kind {
-						let type_lookup = resolve_user_defined_type(udt, found_env, scope_visitor.found_stmt_index.unwrap_or(0));
+						let type_lookup = resolve_user_defined_type(udt, &found_env, scope_visitor.found_stmt_index.unwrap_or(0));
 
 						let completions = if let Ok(type_lookup) = type_lookup {
 							get_completions_from_type(&type_lookup, &types, Some(found_env.phase), false)
@@ -280,15 +278,15 @@ pub fn on_completion(params: lsp_types::CompletionParams) -> CompletionResponse
 					ExprKind::Call { arg_list, callee } => Some((
 						match callee {
 							CalleeKind::Expr(expr) => types.get_expr_type(expr),
-							CalleeKind::SuperCall(method) => resolve_super_method(method, found_env, &types).map_or(types.error(), |(t,_)| t),
+							CalleeKind::SuperCall(method) => resolve_super_method(method, &found_env, &types).map_or(types.error(), |(t,_)| t),
 						}
 						, arg_list)),
 					ExprKind::New(new_expr) => {
 						Some((types.get_expr_type(&new_expr.class), &new_expr.arg_list))
 					}
 					_ => None,
 				}) {
-					let mut completions = get_current_scope_completions(&scope_visitor, &node_to_complete, &preceding_text);
+					let mut completions = get_current_scope_completions(&types, &scope_visitor, &node_to_complete, &preceding_text);
 
 					let arg_list_strings = &callish_expr
 						.1
@@ -328,7 +326,7 @@ pub fn on_completion(params: lsp_types::CompletionParams) -> CompletionResponse
 			}
 
 			// fallback: no special completions, just get stuff from the current scope
-			get_current_scope_completions(&scope_visitor, &node_to_complete, &preceding_text)
+			get_current_scope_completions(&types, &scope_visitor, &node_to_complete, &preceding_text)
 		});
 
 		final_completions = final_completions
@@ -346,6 +344,7 @@ pub fn on_completion(params: lsp_types::CompletionParams) -> CompletionResponse
 
 /// Get symbols in the current scope as completion items
 fn get_current_scope_completions(
+	types: &Types,
 	scope_visitor: &ScopeVisitor,
 	node_to_complete: &Node,
 	preceding_text: &String,
@@ -461,8 +460,7 @@ fn get_current_scope_completions(
 	}
 
 	let found_stmt_index = scope_visitor.found_stmt_index.unwrap_or_default();
-	let found_env = scope_visitor.found_scope.env.borrow();
-	let found_env = found_env.as_ref().unwrap();
+	let found_env = types.get_scope_env(&scope_visitor.found_scope);
 
 	for symbol_data in found_env.symbol_map.iter().filter(|s| {
 		if let StatementIdx::Index(i) = s.1 .0 {

libs/wingc/src/lsp/hover.rs

@@ -43,8 +43,7 @@ impl<'a> HoverVisitor<'a> {
 	/// Try to look up a full path of a symbol in the current scope and if found, render the docs
 	/// associated with the symbol kind. Returns `None` if not found.
 	fn lookup_docs(&mut self, nested_str: &str, property: Option<&Symbol>) -> Option<String> {
-		let current_env = self.current_scope.env.borrow();
-		let current_env = current_env.as_ref()?;
+		let current_env = self.types.get_scope_env(self.current_scope);
 
 		let result = current_env.lookup_nested_str(nested_str, None);
 
@@ -242,12 +241,11 @@ impl<'a> Visit<'a> for HoverVisitor<'a> {
 			ExprKind::Call { arg_list, callee } => {
 				let x = arg_list.named_args.iter().find(|a| a.0.span.contains(&self.position));
 				if let Some((arg_name, ..)) = x {
-					let curr_env = self.current_scope.env.borrow();
-					let env = curr_env.as_ref().expect("an env");
+					let env = self.types.get_scope_env(self.current_scope);
 					// we need to get the struct type from the callee
 					let callee_type = match callee {
 						CalleeKind::Expr(expr) => self.types.get_expr_type(expr),
-						CalleeKind::SuperCall(method) => resolve_super_method(method, env, &self.types)
+						CalleeKind::SuperCall(method) => resolve_super_method(method, &env, &self.types)
 							.ok()
 							.map_or(self.types.error(), |t| t.0),
 					};

libs/wingc/src/lsp/signature.rs

@@ -10,7 +10,7 @@ use crate::lsp::sync::PROJECT_DATA;
 use crate::lsp::sync::WING_TYPES;
 
 use crate::type_check::symbol_env::SymbolEnvRef;
-use crate::type_check::{resolve_super_method, resolve_user_defined_type, CLASS_INIT_NAME};
+use crate::type_check::{resolve_super_method, resolve_user_defined_type, Types, CLASS_INIT_NAME};
 use crate::visit::{visit_expr, visit_scope, Visit};
 use crate::wasm_util::{ptr_to_string, string_to_combined_ptr, WASM_RETURN_ERROR};
 
@@ -37,7 +37,7 @@ pub fn on_signature_help(params: lsp_types::SignatureHelpParams) -> Option<Signa
 			let file = uri.to_file_path().ok().expect("LSP only works on real filesystems");
 			let root_scope = &project_data.asts.get(&file).unwrap();
 
-			let mut scope_visitor = ScopeVisitor::new(params.text_document_position_params.position);
+			let mut scope_visitor = ScopeVisitor::new(&types, params.text_document_position_params.position);
 			scope_visitor.visit_scope(root_scope);
 			let expr = scope_visitor.call_expr?;
 			let env = scope_visitor.call_env?;
@@ -52,7 +52,7 @@ pub fn on_signature_help(params: lsp_types::SignatureHelpParams) -> Option<Signa
 						return None;
 					};
 
-					let Some(t) = resolve_user_defined_type(&udt, root_scope.env.borrow().as_ref()?, 0).ok() else {
+					let Some(t) = resolve_user_defined_type(&udt, &types.get_scope_env(&root_scope), 0).ok() else {
 						return None;
 					};
 
@@ -181,6 +181,7 @@ pub fn on_signature_help(params: lsp_types::SignatureHelpParams) -> Option<Signa
 /// This visitor is used to find the scope
 /// and relevant expression that contains a given location.
 pub struct ScopeVisitor<'a> {
+	types: &'a Types,
 	/// The target location we're looking for
 	pub location: Position,
 	/// The nearest expression before (or containing) the target location
@@ -192,8 +193,9 @@ pub struct ScopeVisitor<'a> {
 }
 
 impl<'a> ScopeVisitor<'a> {
-	pub fn new(location: Position) -> Self {
+	pub fn new(types: &'a Types, location: Position) -> Self {
 		Self {
+			types,
 			location,
 			call_expr: None,
 			call_env: None,
@@ -222,7 +224,8 @@ impl<'a> Visit<'a> for ScopeVisitor<'a> {
 	}
 
 	fn visit_scope(&mut self, node: &'a crate::ast::Scope) {
-		self.curr_env.push(node.env.borrow().as_ref().unwrap().get_ref());
+		let env = self.types.get_scope_env(node);
+		self.curr_env.push(env);
 		visit_scope(self, node);
 		self.curr_env.pop();
 	}

libs/wingc/src/lsp/sync.rs

@@ -14,7 +14,6 @@ use crate::fold::Fold;
 use crate::jsify::JSifier;
 use crate::lifting::LiftVisitor;
 use crate::parser::parse_wing_project;
-use crate::reset::ScopeResetter;
 use crate::type_check;
 use crate::type_check::jsii_importer::JsiiImportSpec;
 use crate::type_check_assert::TypeCheckAssert;
@@ -155,10 +154,7 @@ fn partial_compile(
 
 	// Reset all type information
 	types.reset_expr_types();
-	let mut scope_resetter = ScopeResetter::new();
-	for scope in project_data.asts.values() {
-		scope_resetter.reset_scopes(scope);
-	}
+	types.reset_scope_envs();
 
 	// -- TYPECHECKING PHASE --
 	let mut jsii_imports = vec![];
@@ -167,7 +163,6 @@ fn partial_compile(
 	// Wing files, then move on to files that depend on those, etc.)
 	for file in &topo_sorted_files {
 		let mut scope = project_data.asts.get_mut(file).expect("matching AST not found");
-		scope.env.borrow_mut().take();
 		type_check(&mut scope, &mut types, &file, jsii_types, &mut jsii_imports);
 
 		// Validate the type checker didn't miss anything - see `TypeCheckAssert` for details