ruby-rbs crate: Add declaration tables to Environment

rust/ruby-rbs/Cargo.toml

@@ -17,6 +17,7 @@ include = [
 [dependencies]
 ruby-rbs-sys = { version = "0.3", path = "../ruby-rbs-sys" }
 xxhash-rust = { version = "0.8", features = ["xxh3"] }
+indexmap = "2"
 
 [build-dependencies]
 serde = { version = "1.0", features = ["derive"] }

rust/ruby-rbs/src/ast/declarations.rs

@@ -3,7 +3,7 @@ use crate::ast::comment::Comment;
 use crate::ast::location::{
     AliasDeclarationLocation, ClassDeclarationLocation, ClassSuperLocation,
     ConstantDeclarationLocation, GlobalDeclarationLocation, InterfaceDeclarationLocation,
-    ModuleDeclarationLocation, ModuleSelfLocation, TypeAliasDeclarationLocation,
+    LocationRange, ModuleDeclarationLocation, ModuleSelfLocation, TypeAliasDeclarationLocation,
 };
 use crate::ast::members::Member;
 use crate::ast::type_param::TypeParam;
@@ -22,18 +22,54 @@ pub enum Declaration {
     ModuleAlias(ModuleAliasDeclaration),
 }
 
+impl Declaration {
+    #[must_use]
+    pub(crate) fn location_range(&self) -> Option<LocationRange> {
+        match self {
+            Declaration::Class(d) => d.location.as_ref().map(|l| l.range),
+            Declaration::Module(d) => d.location.as_ref().map(|l| l.range),
+            Declaration::Interface(d) => d.location.as_ref().map(|l| l.range),
+            Declaration::Constant(d) => d.location.as_ref().map(|l| l.range),
+            Declaration::Global(d) => d.location.as_ref().map(|l| l.range),
+            Declaration::TypeAlias(d) => d.location.as_ref().map(|l| l.range),
+            Declaration::ClassAlias(d) => d.location.as_ref().map(|l| l.range),
+            Declaration::ModuleAlias(d) => d.location.as_ref().map(|l| l.range),
+        }
+    }
+}
+
 #[derive(Clone, Debug, Eq, PartialEq, Hash)]
 pub enum ClassMember {
     Member(Member),
     Declaration(Declaration),
 }
 
+impl ClassMember {
+    #[must_use]
+    pub fn as_declaration(&self) -> Option<&Declaration> {
+        match self {
+            Self::Declaration(d) => Some(d),
+            Self::Member(_) => None,
+        }
+    }
+}
+
 #[derive(Clone, Debug, Eq, PartialEq, Hash)]
 pub enum ModuleMember {
     Member(Member),
     Declaration(Declaration),
 }
 
+impl ModuleMember {
+    #[must_use]
+    pub fn as_declaration(&self) -> Option<&Declaration> {
+        match self {
+            Self::Declaration(d) => Some(d),
+            Self::Member(_) => None,
+        }
+    }
+}
+
 #[derive(Clone, Debug, Eq, PartialEq, Hash)]
 pub struct ClassSuper {
     pub name: TypeName,

rust/ruby-rbs/src/environment/decl.rs

@@ -0,0 +1,76 @@
+use crate::ast::Declaration;
+use crate::ids::TypeName;
+
+use super::Environment;
+
+#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
+pub struct DeclId(u32);
+
+#[derive(Copy, Clone, Debug)]
+pub(super) struct DeclNode {
+    pub(super) source: u32,
+    pub(super) parent: Option<DeclId>,
+    pub(super) index: u32,
+    pub(super) context: Context,
+}
+
+#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
+pub struct Context(Option<u32>);
+
+impl Context {
+    #[must_use]
+    pub const fn root() -> Self {
+        Self(None)
+    }
+}
+
+#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
+pub(super) struct ContextNode {
+    parent: Context,
+    name: TypeName,
+}
+
+impl Environment {
+    #[must_use]
+    pub fn decl(&self, id: DeclId) -> &Declaration {
+        self.member_decl(self.node(id))
+    }
+
+    #[must_use]
+    pub fn context_of(&self, id: DeclId) -> Context {
+        self.node(id).context
+    }
+
+    pub fn context_names(&self, context: Context) -> impl Iterator<Item = TypeName> + '_ {
+        let node = move |i: u32| &self.contexts[i as usize];
+        std::iter::successors(context.0, move |&i| node(i).parent.0).map(move |i| node(i).name)
+    }
+
+    pub(super) fn intern_context(&mut self, parent: Context, name: TypeName) -> Context {
+        let (index, _) = self.contexts.insert_full(ContextNode { parent, name });
+        Context(Some(u32::try_from(index).expect("too many contexts")))
+    }
+
+    pub(super) fn node(&self, id: DeclId) -> &DeclNode {
+        &self.decl_nodes[id.0 as usize]
+    }
+
+    pub(super) fn push_node(&mut self, node: DeclNode) -> DeclId {
+        let id = DeclId(u32::try_from(self.decl_nodes.len()).expect("too many declarations"));
+        self.decl_nodes.push(node);
+        id
+    }
+
+    pub(super) fn member_decl(&self, node: &DeclNode) -> &Declaration {
+        let index = node.index as usize;
+        let Some(parent) = node.parent else {
+            return &self.sources[node.source as usize].declarations[index];
+        };
+        let member = match self.decl(parent) {
+            Declaration::Class(c) => c.members[index].as_declaration(),
+            Declaration::Module(m) => m.members[index].as_declaration(),
+            _ => unreachable!("only class/module declarations have child declarations"),
+        };
+        member.expect("DeclNode must point at a Declaration")
+    }
+}