Swift: generate raw helpers in synthesized stubs

redsun82 · redsun82 · commit e2d7a6910c4e · 2023-02-16T10:49:21.000+01:00
This will add helpers to get the underlying raw entities or constructor
arguments on stubs for synthesized classes.

For example a schema like:

```
@synth.from_class(A)
class B:
    pass

@synth.on_arguments(base=A, index=int)
class C:
    pass
```

will generate

```
cached
private Raw::A getUnderlyingEntity() { this = Synth::TB(result) }
```
in the `B.qll` stub and
```
cached
private Raw::A getUnderlyingBase() { this = Synth::TC(result, _) }

cached
private int getUnderlyingIndex() { this = Synth::TC(_, result) }
```
in the `C.qll` stub.

As stubs these can be freely changed later on.
diff --git a/swift/.gitignore b/swift/.gitignore
@@ -11,5 +11,6 @@ compile_commands.json
 /.idea
 /cmake*
 
-# VSCode default build directory
+# VSCode default build directory and project directory
 /build
+/.vscode
diff --git a/swift/codegen/generators/qlgen.py b/swift/codegen/generators/qlgen.py
@@ -146,7 +146,7 @@ def get_ql_property(cls: schema.Class, prop: schema.Property, prev_child: str =
     return ql.Property(**args)
 
 
-def get_ql_class(cls: schema.Class):
+def get_ql_class(cls: schema.Class) -> ql.Class:
     pragmas = {k: True for k in cls.pragmas if k.startswith("ql")}
     prev_child = ""
     properties = []
@@ -228,6 +228,10 @@ def format(codeql, files):
         log.debug(line.strip())
 
 
+def _get_path(cls: schema.Class) -> pathlib.Path:
+    return pathlib.Path(cls.group or "", cls.name).with_suffix(".qll")
+
+
 def _get_all_properties(cls: schema.Class, lookup: typing.Dict[str, schema.Class],
                         already_seen: typing.Optional[typing.Set[int]] = None) -> \
         typing.Iterable[typing.Tuple[schema.Class, schema.Property]]:
@@ -283,6 +287,29 @@ def _should_skip_qltest(cls: schema.Class, lookup: typing.Dict[str, schema.Class
         cls, lookup)
 
 
+def _get_stub(cls: schema.Class, base_import: str) -> ql.Stub:
+    if isinstance(cls.ipa, schema.IpaInfo):
+        if cls.ipa.from_class is not None:
+            accessors = [
+                ql.IpaUnderlyingAccessor(
+                    argument="Entity",
+                    type=_to_db_type(cls.ipa.from_class),
+                    constructorparams=["result"]
+                )
+            ]
+        elif cls.ipa.on_arguments is not None:
+            accessors = [
+                ql.IpaUnderlyingAccessor(
+                    argument=inflection.camelize(arg),
+                    type=_to_db_type(type),
+                    constructorparams=["result" if a == arg else "_" for a in cls.ipa.on_arguments]
+                ) for arg, type in cls.ipa.on_arguments.items()
+            ]
+    else:
+        accessors = []
+    return ql.Stub(name=cls.name, base_import=base_import, ipa_accessors=accessors)
+
+
 def generate(opts, renderer):
     input = opts.schema
     out = opts.ql_output
@@ -323,10 +350,13 @@ def generate(opts, renderer):
             qll = out / c.path.with_suffix(".qll")
             c.imports = [imports[t] for t in get_classes_used_by(c)]
             renderer.render(c, qll)
-            stub_file = stub_out / c.path.with_suffix(".qll")
+
+        for c in data.classes.values():
+            path = _get_path(c)
+            stub_file = stub_out / path
             if not renderer.is_customized_stub(stub_file):
-                stub = ql.Stub(name=c.name, base_import=get_import(qll, opts.swift_dir))
-                renderer.render(stub, stub_file)
+                base_import = get_import(out / path, opts.swift_dir)
+                renderer.render(_get_stub(c, base_import), stub_file)
 
         # for example path/to/elements -> path/to/elements.qll
         renderer.render(ql.ImportList([i for name, i in imports.items() if not classes[name].ql_internal]),
diff --git a/swift/codegen/lib/ql.py b/swift/codegen/lib/ql.py
@@ -134,12 +134,29 @@ def has_doc(self) -> bool:
         return bool(self.doc) or self.ql_internal
 
 
+@dataclass
+class IpaUnderlyingAccessor:
+    argument: str
+    type: str
+    constructorparams: List[Param]
+
+    def __post_init__(self):
+        if self.constructorparams:
+            self.constructorparams = [Param(x) for x in self.constructorparams]
+            self.constructorparams[0].first = True
+
+
 @dataclass
 class Stub:
     template: ClassVar = 'ql_stub'
 
     name: str
     base_import: str
+    ipa_accessors: List[IpaUnderlyingAccessor] = field(default_factory=list)
+
+    @property
+    def has_ipa_accessors(self) -> bool:
+        return bool(self.ipa_accessors)
 
 
 @dataclass
diff --git a/swift/codegen/templates/ql_stub.mustache b/swift/codegen/templates/ql_stub.mustache
@@ -1,9 +1,18 @@
 // generated by {{generator}}, remove this comment if you wish to edit this file
 private import {{base_import}}
+{{#has_ipa_accessors}}
+private import codeql.swift.generated.Raw
+private import codeql.swift.generated.Synth
+{{/has_ipa_accessors}}
 
 {{#ql_internal}}
 /**
  * INTERNAL: Do not use.
  */
 {{/ql_internal}}
-class {{name}} extends Generated::{{name}} {}
+class {{name}} extends Generated::{{name}} {
+    {{#ipa_accessors}}
+    private
+    cached {{type}} getUnderlying{{argument}}() { this = Synth::T{{name}}({{#constructorparams}}{{^first}},{{/first}}{{param}}{{/constructorparams}})}
+    {{/ipa_accessors}}
+}
diff --git a/swift/codegen/test/test_ql.py b/swift/codegen/test/test_ql.py
@@ -149,5 +149,12 @@ def test_class_without_description():
     assert prop.has_description is False
 
 
+def test_ipa_accessor_has_first_constructor_param_marked():
+    params = ["a", "b", "c"]
+    x = ql.IpaUnderlyingAccessor("foo", "bar", params)
+    assert x.constructorparams[0].first
+    assert [p.param for p in x.constructorparams] == params
+
+
 if __name__ == '__main__':
     sys.exit(pytest.main([__file__] + sys.argv[1:]))
diff --git a/swift/codegen/test/test_qlgen.py b/swift/codegen/test/test_qlgen.py
@@ -816,5 +816,29 @@ def test_property_on_class_with_default_doc_name(generate_classes):
     }
 
 
+def test_stub_on_class_with_ipa_from_class(generate_classes):
+    assert generate_classes([
+        schema.Class("MyObject", ipa=schema.IpaInfo(from_class="A")),
+    ]) == {
+        "MyObject.qll": (ql.Stub(name="MyObject", base_import=gen_import_prefix + "MyObject", ipa_accessors=[
+            ql.IpaUnderlyingAccessor(argument="Entity", type="Raw::A", constructorparams=["result"]),
+        ]),
+            ql.Class(name="MyObject", final=True, ipa=True)),
+    }
+
+
+def test_stub_on_class_with_ipa_on_arguments(generate_classes):
+    assert generate_classes([
+        schema.Class("MyObject", ipa=schema.IpaInfo(on_arguments={"base": "A", "index": "int", "label": "string"})),
+    ]) == {
+        "MyObject.qll": (ql.Stub(name="MyObject", base_import=gen_import_prefix + "MyObject", ipa_accessors=[
+            ql.IpaUnderlyingAccessor(argument="Base", type="Raw::A", constructorparams=["result", "_", "_"]),
+            ql.IpaUnderlyingAccessor(argument="Index", type="int", constructorparams=["_", "result", "_"]),
+            ql.IpaUnderlyingAccessor(argument="Label", type="string", constructorparams=["_", "_", "result"]),
+        ]),
+            ql.Class(name="MyObject", final=True, ipa=True)),
+    }
+
+
 if __name__ == '__main__':
     sys.exit(pytest.main([__file__] + sys.argv[1:]))
diff --git a/swift/ql/.generated.list b/swift/ql/.generated.list
