From 543d27c2f0e43113fb22740a90946eeb52e57bb5 Mon Sep 17 00:00:00 2001
From: Joe Savona <joesavona@fb.com>
Date: Thu, 17 Aug 2023 10:02:56 -0700
Subject: [PATCH] [rust][sema] Support classes

Adds name resolution support for class declarations and expressions. Mostly this
involves _not_ visiting some `Identifier` nodes that don't actually represent
variables. For example, method names don't introduce new variables (but if they
are computed, they may _refer_ to variables).
---
 .../forget_estree/src/generated_extensions.rs | 44 ++++++++++-
 .../forget_semantic_analysis/src/analyzer.rs  | 76 ++++++++++++++++++-
 2 files changed, 114 insertions(+), 6 deletions(-)

diff --git a/compiler/forget/crates/forget_estree/src/generated_extensions.rs b/compiler/forget/crates/forget_estree/src/generated_extensions.rs
index 007aab7c9d..f746c02f5c 100644
--- a/compiler/forget/crates/forget_estree/src/generated_extensions.rs
+++ b/compiler/forget/crates/forget_estree/src/generated_extensions.rs
@@ -1,8 +1,8 @@
 // Manual extensions to generated types
 use crate::{
-    ArrowFunctionExpression, Function, FunctionDeclaration, FunctionExpression,
-    ImportDeclarationSpecifier, JSXElementName, JSXMemberExpression,
-    JSXMemberExpressionOrIdentifier, Pattern, SourceRange, SourceType,
+    ArrowFunctionExpression, Class, ClassDeclaration, ClassExpression, Function,
+    FunctionDeclaration, FunctionExpression, ImportDeclarationSpecifier, JSXElementName,
+    JSXMemberExpression, JSXMemberExpressionOrIdentifier, Pattern, SourceRange, SourceType,
 };
 
 /// Sentinel trait to distinguish AST *node* types
@@ -102,3 +102,41 @@ impl IntoFunction for Function {
         self
     }
 }
+
+pub trait IntoClass: ESTreeNode {
+    fn class(&self) -> &Class;
+
+    fn into_class(self) -> Class;
+}
+
+impl IntoClass for ClassDeclaration {
+    fn class(&self) -> &Class {
+        &self.class
+    }
+
+    fn into_class(self) -> Class {
+        self.class
+    }
+}
+
+impl IntoClass for ClassExpression {
+    fn class(&self) -> &Class {
+        &self.class
+    }
+
+    fn into_class(self) -> Class {
+        self.class
+    }
+}
+
+impl ESTreeNode for Class {}
+
+impl IntoClass for Class {
+    fn class(&self) -> &Class {
+        self
+    }
+
+    fn into_class(self) -> Class {
+        self
+    }
+}
diff --git a/compiler/forget/crates/forget_semantic_analysis/src/analyzer.rs b/compiler/forget/crates/forget_semantic_analysis/src/analyzer.rs
index c218431af4..75d8bf4204 100644
--- a/compiler/forget/crates/forget_semantic_analysis/src/analyzer.rs
+++ b/compiler/forget/crates/forget_semantic_analysis/src/analyzer.rs
@@ -1,9 +1,9 @@
 use forget_diagnostics::Diagnostic;
 use forget_estree::{
     AssignmentOperator, AssignmentPropertyOrRestElement, AssignmentTarget, Expression,
-    ExpressionOrSuper, ForInInit, ForInit, FunctionBody, Identifier, ImportDeclarationSpecifier,
-    IntoFunction, JSXElementName, Pattern, Program, SourceRange, Statement,
-    VariableDeclarationKind, Visitor,
+    ExpressionOrPrivateIdentifier, ExpressionOrSuper, ForInInit, ForInit, FunctionBody, Identifier,
+    ImportDeclarationSpecifier, IntoClass, IntoFunction, JSXElementName, Pattern, Program,
+    SourceRange, Statement, VariableDeclarationKind, Visitor,
 };
 
 use crate::{
@@ -186,6 +186,18 @@ impl Analyzer {
             .insert(AstNode::from(function), scope);
     }
 
+    fn visit_class<T: IntoClass>(&mut self, node: &T) {
+        let class = node.class();
+        let scope = self.enter(ScopeKind::Class, |visitor| {
+            if let Some(super_class) = &class.super_class {
+                visitor.visit_expression(super_class);
+            }
+
+            visitor.visit_class_body(&class.body);
+        });
+        self.manager.node_scopes.insert(AstNode::from(class), scope);
+    }
+
     fn visit_reference_identifier(
         &mut self,
         name: &str,
@@ -337,6 +349,64 @@ impl Visitor for Analyzer {
         }
     }
 
+    fn visit_class_declaration(&mut self, ast: &forget_estree::ClassDeclaration) {
+        if let Some(id) = &ast.class.id {
+            let declaration = self.manager.add_declaration(
+                self.current,
+                id.name.clone(),
+                DeclarationKind::Class,
+                id.range,
+            );
+            self.manager
+                .node_declarations
+                .insert(AstNode::from(id), declaration);
+        }
+        Analyzer::visit_class(self, ast);
+    }
+
+    fn visit_class_property(&mut self, ast: &forget_estree::ClassProperty) {
+        // Static (non-computed) property names do not introduce a new identifier
+        // into any scope
+        if ast.is_computed {
+            self.visit_expression(&ast.key);
+        }
+        if let Some(value) = &ast.value {
+            self.visit_expression(value)
+        }
+    }
+
+    fn visit_class_private_property(&mut self, ast: &forget_estree::ClassPrivateProperty) {
+        // Static (non-computed) property names do not introduce a new identifier
+        // into any scope
+        match &ast.key {
+            ExpressionOrPrivateIdentifier::Expression(key) => {
+                self.visit_expression(key);
+            }
+            ExpressionOrPrivateIdentifier::PrivateIdentifier(_)
+            | ExpressionOrPrivateIdentifier::PrivateName(_) => { /* no-op */ }
+        }
+        if let Some(value) = &ast.value {
+            self.visit_expression(value)
+        }
+    }
+
+    fn visit_static_block(&mut self, ast: &forget_estree::StaticBlock) {
+        self.enter(ScopeKind::StaticBlock, |visitor| {
+            for statement in &ast.body {
+                visitor.visit_statement(statement);
+            }
+        });
+    }
+
+    fn visit_method_definition(&mut self, ast: &forget_estree::MethodDefinition) {
+        // Static (non-computed) method names do not introduce a new identifier
+        // into any scope
+        if ast.is_computed {
+            self.visit_expression(&ast.key)
+        }
+        self.visit_function_expression(&ast.value);
+    }
+
     fn visit_function_declaration(&mut self, ast: &forget_estree::FunctionDeclaration) {
         if let Some(id) = &ast.function.id {
             let declaration = self.manager.add_declaration(