Merge pull request #17918 from hvitved/rust/cfg-codegen

Rust: Add (auto-generated) CFG node wrapper classes

Author: hvitved

misc/codegen/codegen.py

@@ -54,6 +54,8 @@ def _parse_args() -> argparse.Namespace:
                             "generated qll file importing every class file"),
         p.add_argument("--ql-test-output",
                        help="output directory for QL generated extractor test files"),
+        p.add_argument("--ql-cfg-output",
+                       help="output directory for QL CFG layer (optional)."),
         p.add_argument("--cpp-output",
                        help="output directory for generated C++ files, required if trap or cpp is provided to "
                             "--generate"),

misc/codegen/generators/qlgen.py

@@ -160,6 +160,8 @@ def get_ql_class(cls: schema.Class, lookup: typing.Dict[str, schema.Class]) -> q
         prop = get_ql_property(cls, p, lookup, prev_child)
         if prop.is_child:
             prev_child = prop.singular
+            if prop.type in lookup and lookup[prop.type].cfg:
+                prop.cfg = True
         properties.append(prop)
     return ql.Class(
         name=cls.name,
@@ -171,6 +173,16 @@ def get_ql_class(cls: schema.Class, lookup: typing.Dict[str, schema.Class]) -> q
         doc=cls.doc,
         hideable="ql_hideable" in cls.pragmas,
         internal="ql_internal" in cls.pragmas,
+        cfg=cls.cfg,
+    )
+
+
+def get_ql_cfg_class(cls: schema.Class, lookup: typing.Dict[str, ql.Class]) -> ql.CfgClass:
+    return ql.CfgClass(
+        name=cls.name,
+        bases=[base for base in cls.bases if lookup[base.base].cfg],
+        properties=cls.properties,
+        doc=cls.doc
     )
 
 
@@ -361,6 +373,7 @@ def generate(opts, renderer):
     input = opts.schema
     out = opts.ql_output
     stub_out = opts.ql_stub_output
+    cfg_out = opts.ql_cfg_output
     test_out = opts.ql_test_output
     missing_test_source_filename = "MISSING_SOURCE.txt"
     include_file = stub_out.with_suffix(".qll")
@@ -385,6 +398,7 @@ def generate(opts, renderer):
     imports = {}
     imports_impl = {}
     classes_used_by = {}
+    cfg_classes = []
     generated_import_prefix = get_import(out, opts.root_dir)
     registry = opts.generated_registry or pathlib.Path(
         os.path.commonpath((out, stub_out, test_out)), ".generated.list")
@@ -402,6 +416,8 @@ def generate(opts, renderer):
             imports[c.name] = path
             path_impl = get_import(stub_out / c.dir / "internal" / c.name, opts.root_dir)
             imports_impl[c.name + "Impl"] = path_impl + "Impl"
+            if c.cfg:
+                cfg_classes.append(get_ql_cfg_class(c, classes))
 
         for c in classes.values():
             qll = out / c.path.with_suffix(".qll")
@@ -411,6 +427,14 @@ def generate(opts, renderer):
             c.import_prefix = generated_import_prefix
             renderer.render(c, qll)
 
+        if cfg_out:
+            cfg_classes_val = ql.CfgClasses(
+                include_file_import=get_import(include_file, opts.root_dir),
+                classes=cfg_classes
+            )
+            cfg_qll = cfg_out / "CfgNodes.qll"
+            renderer.render(cfg_classes_val, cfg_qll)
+
         for c in data.classes.values():
             path = _get_path(c)
             path_impl = _get_path_impl(c)

misc/codegen/lib/ql.py

@@ -45,6 +45,7 @@ class Property:
     synth: bool = False
     type_is_hideable: bool = False
     internal: bool = False
+    cfg: bool = False
 
     def __post_init__(self):
         if self.tableparams:
@@ -110,6 +111,7 @@ class Class:
     internal: bool = False
     doc: List[str] = field(default_factory=list)
     hideable: bool = False
+    cfg: bool = False
 
     def __post_init__(self):
         def get_bases(bases): return [Base(str(b), str(prev)) for b, prev in zip(bases, itertools.chain([""], bases))]
@@ -333,3 +335,18 @@ class ConstructorStub:
 
         cls: "Synth.FinalClass"
         import_prefix: str
+
+
+@dataclass
+class CfgClass:
+    name: str
+    bases: List[Base] = field(default_factory=list)
+    properties: List[Property] = field(default_factory=list)
+    doc: List[str] = field(default_factory=list)
+
+
+@dataclass
+class CfgClasses:
+    template: ClassVar = 'ql_cfg_nodes'
+    include_file_import: Optional[str] = None
+    classes: List[CfgClass] = field(default_factory=list)

misc/codegen/lib/schema.py

@@ -94,6 +94,7 @@ class Class:
     properties: List[Property] = field(default_factory=list)
     pragmas: List[str] | Dict[str, object] = field(default_factory=dict)
     doc: List[str] = field(default_factory=list)
+    cfg: bool = False
 
     def __post_init__(self):
         if not isinstance(self.pragmas, dict):

misc/codegen/lib/schemadefs.py

@@ -279,7 +279,7 @@ def __or__(self, other: _schema.PropertyModifier):
 drop = object()
 
 
-def annotate(annotated_cls: type, add_bases: _Iterable[type] | None = None, replace_bases: _Dict[type, type] | None = None) -> _Callable[[type], _PropertyAnnotation]:
+def annotate(annotated_cls: type, add_bases: _Iterable[type] | None = None, replace_bases: _Dict[type, type] | None = None, cfg: bool = False) -> _Callable[[type], _PropertyAnnotation]:
     """
     Add or modify schema annotations after a class has been defined previously.
 
@@ -298,6 +298,7 @@ def decorator(cls: type) -> _PropertyAnnotation:
             annotated_cls.__bases__ = tuple(replace_bases.get(b, b) for b in annotated_cls.__bases__)
         if add_bases:
             annotated_cls.__bases__ += tuple(add_bases)
+        annotated_cls.__cfg__ = cfg
         for a in dir(cls):
             if a.startswith(_schema.inheritable_pragma_prefix):
                 setattr(annotated_cls, a, getattr(cls, a))

misc/codegen/loaders/schemaloader.py

@@ -53,6 +53,7 @@ def _get_class(cls: type) -> schema.Class:
                         bases=[b.__name__ for b in cls.__bases__ if b is not object],
                         derived=derived,
                         pragmas=pragmas,
+                        cfg=cls.__cfg__ if hasattr(cls, "__cfg__") else False,
                         # in the following we don't use `getattr` to avoid inheriting
                         properties=[
                             a | _PropertyNamer(n)

misc/codegen/templates/ql_cfg_nodes.mustache

@@ -0,0 +1,199 @@
+// generated by {{generator}}, do not edit
+/**
+ * This module provides generated wrappers around the `CfgNode` type.
+ *
+ * INTERNAL: Do not import directly.
+ */
+
+private import codeql.util.Location
+private import codeql.util.Unit
+private import {{include_file_import}}
+
+/** Provides the input to `MakeCfgNodes` */
+signature module InputSig<LocationSig Loc> {
+  class CfgNode {
+    AstNode getAstNode();
+
+    string toString();
+
+    Loc getLocation();
+  }
+
+  AstNode getDesugared(AstNode n);
+}
+
+/**
+ * Given a `CfgNode` implementation, provides the module `Nodes` that
+ * contains wrappers around `CfgNode` for relevant classes.
+ */
+module MakeCfgNodes<LocationSig Loc, InputSig<Loc> Input> {
+  private import Input
+
+  final private class AstNodeFinal = AstNode;
+
+  final private class CfgNodeFinal = CfgNode;
+
+  /**
+   * INTERNAL: Do not expose.
+   */
+  abstract class ParentAstNode extends AstNodeFinal {
+    /**
+     * Holds if `child` is a (possibly nested) child of this AST node
+     * for which we would like to find a matching CFG child.
+     */
+    abstract predicate relevantChild(AstNode child);
+  }
+
+  /**
+   * INTERNAL: Do not expose.
+   */
+  abstract class ChildMapping extends Unit {
+    /**
+     * Holds if `child` is a (possibly nested) child of AST node `parent`
+     * for which we would like to find a matching CFG child.
+     */
+    final predicate relevantChild(AstNode parent, AstNode child) {
+      parent.(ParentAstNode).relevantChild(child)
+    }
+
+    /**
+     * Holds if there is a control flow path from `cfn` to `cfnChild`, where `cfn`
+     * is a control flow node for this AST node, and `cfnChild` is a control flow
+     * node for `child`.
+     *
+     * This predicate should be implemented at the place where `MakeCfgNodes` is
+     * invoked. Ideally, `MakeCfgNodes` should be a higher-order parameterized
+     * module, but since that is currently not supported, we achieve the "callback"
+     * effect using this `abstract` class instead.
+     */
+    cached
+    abstract predicate hasCfgChild(AstNode parent, AstNode child, CfgNode cfn, CfgNode cfnChild);
+  }
+
+  /** Provides sub classes of `CfgNode`. */
+  module Nodes {
+    {{#classes}}
+    private final class Parent{{name}} extends ParentAstNode, {{name}} {
+      override predicate relevantChild(AstNode child) {
+        none()
+        {{#properties}}
+        {{#cfg}}
+        or
+        child = this.{{getter}}({{#is_indexed}}_{{/is_indexed}})
+        {{/cfg}}
+        {{/properties}}
+      }
+    }
+
+    /**
+    {{#doc}}
+    * {{.}}
+    {{/doc}}
+    */
+    final class {{name}}CfgNode extends CfgNodeFinal{{#bases}}, {{.}}CfgNode{{/bases}} {
+      private {{name}} node;
+
+      {{name}}CfgNode() {
+        node = this.getAstNode()
+      }
+
+      /** Gets the underlying `{{name}}`. */
+      {{name}} get{{name}}() { result = node }
+    
+      {{#properties}}
+      /**
+      * {{>ql_property_doc}}     *
+      {{#description}}
+      * {{.}}
+      {{/description}}
+      {{#internal}}
+      * INTERNAL: Do not use.
+      {{/internal}}
+      */
+      {{type}}{{#cfg}}CfgNode{{/cfg}} {{getter}}({{#is_indexed}}int index{{/is_indexed}}) {
+        {{#cfg}}
+        any(ChildMapping mapping).hasCfgChild(node, node.{{getter}}({{#is_indexed}}index{{/is_indexed}}), this, result)
+        {{/cfg}}
+        {{^cfg}}
+        {{^is_predicate}}result = {{/is_predicate}}node.{{getter}}({{#is_indexed}}index{{/is_indexed}})
+        {{/cfg}}
+      }
+
+      {{#is_optional}}
+      /**
+      * Holds if `{{getter}}({{#is_repeated}}index{{/is_repeated}})` exists.
+      {{#internal}}
+      * INTERNAL: Do not use.
+      {{/internal}}
+      */
+      predicate has{{singular}}({{#is_repeated}}int index{{/is_repeated}}) {
+        exists(this.{{getter}}({{#is_repeated}}index{{/is_repeated}}))
+      }
+      {{/is_optional}}
+      {{#is_indexed}}
+
+      /**
+      * Gets any of the {{doc_plural}}.
+      {{#internal}}
+      * INTERNAL: Do not use.
+      {{/internal}}
+      */
+      {{type}}{{#cfg}}CfgNode{{/cfg}} {{indefinite_getter}}() {
+        result = this.{{getter}}(_)
+      }
+      {{^is_optional}}
+
+      /**
+      * Gets the number of {{doc_plural}}.
+      {{#internal}}
+      * INTERNAL: Do not use.
+      {{/internal}}
+      */
+      int getNumberOf{{plural}}() {
+        result = count(int i | exists(this.{{getter}}(i)))
+      }
+      {{/is_optional}}
+      {{/is_indexed}}
+      {{#is_unordered}}
+      /**
+      * Gets the number of {{doc_plural}}.
+      {{#internal}}
+      * INTERNAL: Do not use.
+      {{/internal}}
+      */
+      int getNumberOf{{plural}}() {
+        result = count(this.{{getter}}())
+      }
+      {{/is_unordered}}
+      {{/properties}}
+    }
+    {{/classes}}
+  }
+
+  module Consistency {
+    private predicate hasCfgNode(AstNode astNode) {
+      astNode = any(CfgNode cfgNode).getAstNode()
+    }
+
+    query predicate missingCfgChild(CfgNode parent, string pred, int i, AstNode child) {
+      none()
+      {{#classes}}
+      {{#properties}}
+      {{#cfg}}
+      or
+      pred = "{{getter}}" and
+      parent = any(Nodes::{{name}}CfgNode cfgNode, {{name}} astNode |
+        astNode = cfgNode.get{{name}}() and
+        child = getDesugared(astNode.{{getter}}({{#is_indexed}}i{{/is_indexed}}))
+        {{^is_indexed}}and i = -1{{/is_indexed}} and
+        hasCfgNode(child) and
+        not child = cfgNode.{{getter}}({{#is_indexed}}i{{/is_indexed}}).getAstNode()
+        |
+        cfgNode
+      )
+      {{/cfg}}
+      {{/properties}}
+      {{/classes}}
+    }
+  }
+}

rust/codegen.conf

@@ -4,6 +4,7 @@
 --dbscheme=ql/lib/rust.dbscheme
 --ql-output=ql/lib/codeql/rust/elements/internal/generated
 --ql-stub-output=ql/lib/codeql/rust/elements
+--ql-cfg-output=ql/lib/codeql/rust/controlflow/internal/generated
 --ql-test-output=ql/test/extractor-tests/generated
 --rust-output=extractor/src/generated
 --script-name=codegen