Rust: introduce typed labels

Author: Paolo Tranquilli

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"{t}TrapLabel"
         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 @@ def _get_class(self, name: str) -> rust.Class:
                 _get_field(c, p)
                 for c, p in _get_properties(cls, self._classmap)
                 if "rust_skip" not in p.pragmas and not p.synth
-            ],
+            ] if not cls.derived else [],
+            ancestors=sorted(set(a.name for a in _get_ancestors(cls, self._classmap))),
             table_name=inflection.tableize(cls.name),
         )
 
     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,8 +110,9 @@ def is_label(self):
 @dataclasses.dataclass
 class Class:
     name: str
-    table_name: str
+    table_name: 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):

misc/codegen/templates/rust_classes.mustache

@@ -2,53 +2,104 @@
 
 #![cfg_attr(any(), rustfmt::skip)]
 
-use crate::trap::{TrapId, TrapEntry};
-use codeql_extractor::trap;
+use crate::trap;
 {{#classes}}
 
+#[derive(Debug, Clone, Copy, Eq, PartialEq, Ord, PartialOrd, Hash)]
+pub struct {{name}}TrapLabel(trap::UntypedLabel);
+
+impl From<trap::UntypedLabel> for {{name}}TrapLabel {
+    fn from(value: trap::UntypedLabel) -> Self {
+        Self(value)
+    }
+}
+
+impl  From<{{name}}TrapLabel> for trap::TrapId<{{name}}> {
+    fn from(value: {{name}}TrapLabel) -> Self {
+        Self::Label(value)
+    }
+}
+
+impl trap::Label for {{name}}TrapLabel {
+    fn as_untyped(&self) -> trap::UntypedLabel {
+        self.0
+    }
+}
+
+impl From<{{name}}TrapLabel> for trap::Arg {
+    fn from(value: {{name}}TrapLabel) -> Self {
+        value.0.into()
+    }
+}
+
+{{#table_name}}
 #[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 class_name() -> &'static str { "{{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: Self::Label, 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("{{table_name}}", 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}}
         }
         {{/is_repeated}}
         {{/fields}}
     }
 }
+{{/table_name}}
+{{^table_name}}
+{{! virtual class, make it unbuildable }}
+pub struct {{name}} {
+    unused: ()
+}
+{{/table_name}}
+
+impl trap::TrapClass for {{name}} {
+    type Label = {{name}}TrapLabel;
+}
+{{/classes}}
+
+// Conversions
+{{#classes}}
+{{#ancestors}}
+impl From<{{name}}TrapLabel> for {{.}}TrapLabel {
+    fn from(value: {{name}}TrapLabel) -> Self {
+        value.0.into()
+    }
+}
+{{/ancestors}}
 {{/classes}}