Merge pull request #17460 from github/redsun82/rust-typed-labels

Rust: introduce typed labels

misc/codegen/generators/rustgen.py

@@ -20,7 +20,7 @@ def _get_type(t: str) -> str:
         case "int":
             return "usize"
         case _ if t[0].isupper():
-            return "trap::Label"
+            return f"trap::Label<{t}>"
         case "boolean":
             assert False, "boolean unsupported"
         case _:
@@ -57,6 +57,15 @@ def _get_properties(
         yield cls, p
 
 
+def _get_ancestors(
+    cls: schema.Class, lookup: dict[str, schema.Class]
+) -> typing.Iterable[schema.Class]:
+    for b in cls.bases:
+        base = lookup[b]
+        yield base
+        yield from _get_ancestors(base, lookup)
+
+
 class Processor:
     def __init__(self, data: schema.Schema):
         self._classmap = data.classes
@@ -69,14 +78,15 @@ class Processor:
                 _get_field(c, p)
                 for c, p in _get_properties(cls, self._classmap)
                 if "rust_skip" not in p.pragmas and not p.synth
-            ],
-            table_name=inflection.tableize(cls.name),
+            ] if not cls.derived else [],
+            ancestors=sorted(set(a.name for a in _get_ancestors(cls, self._classmap))),
+            entry_table=inflection.tableize(cls.name) if not cls.derived else None,
         )
 
     def get_classes(self):
         ret = {"": []}
         for k, cls in self._classmap.items():
-            if not cls.synth and not cls.derived:
+            if not cls.synth:
                 ret.setdefault(cls.group, []).append(self._get_class(cls.name))
         return ret
 

misc/codegen/lib/rust.py

@@ -110,12 +110,19 @@ class Field:
 @dataclasses.dataclass
 class Class:
     name: str
-    table_name: str
+    entry_table: str | None = None
     fields: list[Field] = dataclasses.field(default_factory=list)
+    ancestors: list[str] = dataclasses.field(default_factory=list)
 
     @property
-    def single_field_entries(self):
-        ret = {self.table_name: []}
+    def is_entry(self) -> bool:
+        return bool(self.entry_table)
+
+    @property
+    def single_field_entries(self) -> dict[str, list[dict]]:
+        ret = {}
+        if self.is_entry:
+            ret[self.entry_table] = []
         for f in self.fields:
             if f.is_single:
                 ret.setdefault(f.table_name, []).append(f)

misc/codegen/templates/rust_classes.mustache

@@ -2,48 +2,48 @@
 
 #![cfg_attr(any(), rustfmt::skip)]
 
-use crate::trap::{TrapId, TrapEntry};
-use codeql_extractor::trap;
+use crate::trap;
 {{#classes}}
 
+{{#is_entry}}
 #[derive(Debug)]
 pub struct {{name}} {
-    pub id: TrapId,
+    pub id: trap::TrapId<{{name}}>,
     {{#fields}}
     pub {{field_name}}: {{type}},
     {{/fields}}
 }
 
-impl TrapEntry for {{name}} {
-    fn extract_id(&mut self) -> TrapId {
-        std::mem::replace(&mut self.id, TrapId::Star)
+impl trap::TrapEntry for {{name}} {
+    fn extract_id(&mut self) -> trap::TrapId<Self> {
+        std::mem::replace(&mut self.id, trap::TrapId::Star)
     }
 
-    fn emit(self, id: trap::Label, out: &mut trap::Writer) {
+    fn emit(self, id: trap::Label<Self>, out: &mut trap::Writer) {
         {{#single_field_entries}}
-        out.add_tuple("{{table_name}}", vec![trap::Arg::Label(id){{#fields}}, self.{{field_name}}.into(){{/fields}}]);
+        out.add_tuple("{{entry_table}}", vec![id.into(){{#fields}}, self.{{field_name}}.into(){{/fields}}]);
         {{/single_field_entries}}
         {{#fields}}
         {{#is_predicate}}
         if self.{{field_name}} {
-            out.add_tuple("{{table_name}}", vec![trap::Arg::Label(id)]);
+            out.add_tuple("{{table_name}}", vec![id.into()]);
         }
         {{/is_predicate}}
         {{#is_optional}}
         {{^is_repeated}}
         if let Some(v) = self.{{field_name}} {
-            out.add_tuple("{{table_name}}", vec![trap::Arg::Label(id), v.into()]);
+            out.add_tuple("{{table_name}}", vec![id.into(), v.into()]);
         }
         {{/is_repeated}}
         {{/is_optional}}
         {{#is_repeated}}
         for (i, v) in self.{{field_name}}.into_iter().enumerate() {
             {{^is_optional}}
-            out.add_tuple("{{table_name}}", vec![trap::Arg::Label(id){{^is_unordered}}, i.into(){{/is_unordered}}, v.into()]);
+            out.add_tuple("{{table_name}}", vec![id.into(){{^is_unordered}}, i.into(){{/is_unordered}}, v.into()]);
             {{/is_optional}}
             {{#is_optional}}
             if let Some(v) = v {
-                out.add_tuple("{{table_name}}", vec![trap::Arg::Label(id){{^is_unordered}}, i.into(){{/is_unordered}}, v.into()]);
+                out.add_tuple("{{table_name}}", vec![id.into(){{^is_unordered}}, i.into(){{/is_unordered}}, v.into()]);
             }
             {{/is_optional}}
         }
@@ -51,4 +51,27 @@ impl TrapEntry for {{name}} {
         {{/fields}}
     }
 }
+{{/is_entry}}
+{{^is_entry}}
+{{! virtual class, make it unbuildable }}
+#[derive(Debug)]
+pub struct {{name}} {
+    _unused: ()
+}
+{{/is_entry}}
+
+impl trap::TrapClass for {{name}} {
+    fn class_name() -> &'static str { "{{name}}" }
+}
+{{#ancestors}}
+
+impl From<trap::Label<{{name}}>> for trap::Label<{{.}}> {
+    fn from(value: trap::Label<{{name}}>) -> Self {
+        // SAFETY: this is safe because in the dbscheme {{name}} is a subclass of {{.}}
+        unsafe {
+            Self::from_untyped(value.as_untyped())
+        }
+    }
+}
+{{/ancestors}}
 {{/classes}}

rust/extractor/src/generated/.generated.list

@@ -1,2 +1,2 @@
 mod.rs 4bcb9def847469aae9d8649461546b7c21ec97cf6e63d3cf394e339915ce65d7 4bcb9def847469aae9d8649461546b7c21ec97cf6e63d3cf394e339915ce65d7
-top.rs d80ccb71adc47be7dfa395b630da63cb85647886ebc05ca8d53303ebc0c83f76 d80ccb71adc47be7dfa395b630da63cb85647886ebc05ca8d53303ebc0c83f76
+top.rs 4bf2a559a91f46a8fadc7de0c3636d0fc9add8de5b97ae289ade79c1d02a8abd 4bf2a559a91f46a8fadc7de0c3636d0fc9add8de5b97ae289ade79c1d02a8abd

rust/extractor/src/trap.rs

@@ -5,14 +5,19 @@ use log::debug;
 use ra_ap_ide_db::line_index::LineCol;
 use std::ffi::OsString;
 use std::fmt::Debug;
+use std::hash::Hash;
+use std::marker::PhantomData;
 use std::path::{Path, PathBuf};
 
+pub use trap::Label as UntypedLabel;
+pub use trap::Writer;
+
 //TODO: typed labels
 pub trait AsTrapKeyPart {
     fn as_key_part(&self) -> String;
 }
 
-impl AsTrapKeyPart for trap::Label {
+impl AsTrapKeyPart for UntypedLabel {
     fn as_key_part(&self) -> String {
         format!("{{{}}}", self)
     }
@@ -30,28 +35,37 @@ impl AsTrapKeyPart for &str {
     }
 }
 
-#[derive(Debug, Clone)]
-pub enum TrapId {
-    Star,
-    Key(String),
-    Label(trap::Label),
+pub trait TrapClass {
+    fn class_name() -> &'static str;
 }
 
-impl From<String> for TrapId {
+pub trait TrapEntry: Debug + Sized + TrapClass {
+    fn extract_id(&mut self) -> TrapId<Self>;
+    fn emit(self, id: Label<Self>, out: &mut Writer);
+}
+
+#[derive(Debug, Clone)]
+pub enum TrapId<T: TrapEntry> {
+    Star,
+    Key(String),
+    Label(Label<T>),
+}
+
+impl<T: TrapEntry> From<String> for TrapId<T> {
     fn from(value: String) -> Self {
         TrapId::Key(value)
     }
 }
 
-impl From<&str> for TrapId {
+impl<T: TrapEntry> From<&str> for TrapId<T> {
     fn from(value: &str) -> Self {
         TrapId::Key(value.into())
     }
 }
 
-impl From<trap::Label> for TrapId {
-    fn from(value: trap::Label) -> Self {
-        TrapId::Label(value)
+impl<T: TrapEntry> From<Label<T>> for TrapId<T> {
+    fn from(value: Label<T>) -> Self {
+        Self::Label(value)
     }
 }
 
@@ -66,22 +80,61 @@ macro_rules! trap_key {
     }};
 }
 
-pub trait TrapEntry: std::fmt::Debug {
-    fn extract_id(&mut self) -> TrapId;
-    fn emit(self, id: trap::Label, out: &mut trap::Writer);
+#[derive(Debug, Ord, PartialOrd, Eq, PartialEq, Hash)]
+pub struct Label<T: TrapClass> {
+    untyped: UntypedLabel,
+    phantom: PhantomData<T>, // otherwise Rust wants `T` to be used
+}
+
+// not deriving `Clone` and `Copy` because they require `T: Clone` and `T: Copy` respectively,
+// even if `T` is not actually part of the fields.
+// see https://github.com/rust-lang/rust/issues/108894
+impl<T: TrapClass> Clone for Label<T> {
+    fn clone(&self) -> Self {
+        *self
+    }
+}
+
+impl<T: TrapClass> Copy for Label<T> {}
+
+impl<T: TrapClass> Label<T> {
+    pub fn as_untyped(&self) -> UntypedLabel {
+        self.untyped
+    }
+
+    /// # Safety
+    /// The user must make sure the label respects TRAP typing
+    pub unsafe fn from_untyped(untyped: UntypedLabel) -> Self {
+        Self {
+            untyped,
+            phantom: PhantomData,
+        }
+    }
+}
+
+impl<T: TrapClass> AsTrapKeyPart for Label<T> {
+    fn as_key_part(&self) -> String {
+        self.as_untyped().as_key_part()
+    }
+}
+
+impl<T: TrapClass> From<Label<T>> for trap::Arg {
+    fn from(value: Label<T>) -> Self {
+        trap::Arg::Label(value.as_untyped())
+    }
 }
 
 pub struct TrapFile {
     path: PathBuf,
-    writer: trap::Writer,
+    writer: Writer,
     compression: Compression,
 }
 
 impl TrapFile {
-    pub fn emit_location(
+    pub fn emit_location<E: TrapClass>(
         &mut self,
-        file_label: trap::Label,
-        entity_label: trap::Label,
+        file_label: UntypedLabel,
+        entity_label: Label<E>,
         start: LineCol,
         end: LineCol,
     ) {
@@ -101,10 +154,7 @@ impl TrapFile {
         );
         self.writer.add_tuple(
             "locatable_locations",
-            vec![
-                trap::Arg::Label(entity_label),
-                trap::Arg::Label(location_label),
-            ],
+            vec![entity_label.into(), location_label.into()],
         );
     }
 
@@ -112,15 +162,26 @@ impl TrapFile {
         extractor::populate_file(&mut self.writer, absolute_path)
     }
 
-    pub fn label(&mut self, id: TrapId) -> trap::Label {
+    pub fn label<T: TrapEntry>(&mut self, id: TrapId<T>) -> Label<T> {
         match id {
-            TrapId::Star => self.writer.fresh_id(),
-            TrapId::Key(s) => self.writer.global_id(&s).0,
+            TrapId::Star => {
+                let untyped = self.writer.fresh_id();
+                // SAFETY: a `*` trap id is always safe for typing
+                unsafe { Label::from_untyped(untyped) }
+            }
+            TrapId::Key(s) => {
+                let untyped = self
+                    .writer
+                    .global_id(&format!("{},{}", T::class_name(), s))
+                    .0;
+                // SAFETY: using type names as prefixes avoids labels having a conflicting type
+                unsafe { Label::from_untyped(untyped) }
+            }
             TrapId::Label(l) => l,
         }
     }
 
-    pub fn emit<T: TrapEntry>(&mut self, mut e: T) -> trap::Label {
+    pub fn emit<T: TrapEntry>(&mut self, mut e: T) -> Label<T> {
         let label = self.label(e.extract_id());
         e.emit(label, &mut self.writer);
         label

rust/schema.py

@@ -984,7 +984,7 @@ class IdentPat(Pat):
 class TupleStructPat(Pat):
     """
     A tuple struct pattern. For example:
-    ``` 
+    ```
     match x {
         Tuple("a", 1, 2, 3) => "great",
         Tuple(.., 3) => "fine",