Swift: first IPA layer

Author: redsun82

swift/codegen/generators/qlgen.py

@@ -80,6 +80,28 @@ def get_ql_class(cls: schema.Class):
     )
 
 
+def _to_db_type(x: str) -> str:
+    if x[0].isupper():
+        return "@" + inflection.underscore(x)
+    return x
+
+
+_final_db_class_lookup = {}
+
+
+def get_ql_ipa_class(cls: schema.Class):
+    if cls.derived:
+        return ql.Ipa.NonFinalClass(name=cls.name, derived=sorted(cls.derived))
+    if cls.ipa and cls.ipa.from_class:
+        source = cls.ipa.from_class
+        _final_db_class_lookup.setdefault(source, ql.Ipa.FinalClassDb(source)).subtract_type(cls.name)
+        return ql.Ipa.FinalClassIpaFrom(name=cls.name, type=_to_db_type(source))
+    if cls.ipa and cls.ipa.on_arguments:
+        return ql.Ipa.FinalClassIpaOn(name=cls.name,
+                                      params=[ql.Ipa.Param(k, _to_db_type(v)) for k, v in cls.ipa.on_arguments.items()])
+    return _final_db_class_lookup.setdefault(cls.name, ql.Ipa.FinalClassDb(cls.name))
+
+
 def get_import(file: pathlib.Path, swift_dir: pathlib.Path):
     stem = file.relative_to(swift_dir / "ql/lib").with_suffix("")
     return str(stem).replace("/", ".")
@@ -96,7 +118,9 @@ def get_classes_used_by(cls: ql.Class):
     return sorted(set(t for t in get_types_used_by(cls) if t[0].isupper()))
 
 
-_generated_stub_re = re.compile(r"private import .*\n\nclass \w+ extends \w+ \{[ \n]\}", re.MULTILINE)
+_generated_stub_re = re.compile(r"\n*private import .*\n+class \w+ extends \w+ \{[ \n]?\}"
+                                "|"
+                                r"\n*predicate construct\w+\(.*?\) \{ none\(\) \}", re.MULTILINE)
 
 
 def _is_generated_stub(file):
@@ -112,6 +136,7 @@ def _is_generated_stub(file):
         line_threshold = 5
         first_lines = list(itertools.islice(contents, line_threshold))
         if len(first_lines) == line_threshold or not _generated_stub_re.match("".join(first_lines)):
+            print("".join(first_lines))
             raise ModifiedStubMarkedAsGeneratedError(
                 f"{file.name} stub was modified but is still marked as generated")
         return True
@@ -149,12 +174,14 @@ def _get_all_properties_to_be_tested(cls: ql.Class, lookup: typing.Dict[str, ql.
             yield ql.PropertyForTest(p.getter, p.type, p.is_single, p.is_predicate, p.is_repeated)
 
 
+def _partition_iter(x, pred):
+    x1, x2 = itertools.tee(x)
+    return filter(pred, x1), itertools.filterfalse(pred, x2)
+
+
 def _partition(l, pred):
     """ partitions a list according to boolean predicate """
-    res = ([], [])
-    for x in l:
-        res[not pred(x)].append(x)
-    return res
+    return map(list, _partition_iter(l, pred))
 
 
 def _is_in_qltest_collapsed_hierachy(cls: ql.Class, lookup: typing.Dict[str, ql.Class]):
@@ -184,10 +211,10 @@ def generate(opts, renderer):
     existing |= {q for q in test_out.rglob(missing_test_source_filename)}
 
     data = schema.load(input)
+    data.classes.sort(key=lambda cls: (cls.dir, cls.name))
 
     classes = [get_ql_class(cls) for cls in data.classes]
     lookup = {cls.name: cls for cls in classes}
-    classes.sort(key=lambda cls: (cls.dir, cls.name))
     imports = {}
 
     for c in classes:
@@ -228,6 +255,24 @@ def generate(opts, renderer):
             renderer.render(ql.PropertyTester(class_name=c.name,
                                               property=p), test_dir / f"{c.name}_{p.getter}.ql")
 
+    final_ipa_types = []
+    non_final_ipa_types = []
+    constructor_imports = []
+    for cls in data.classes:
+        ipa_type = get_ql_ipa_class(cls)
+        if ipa_type.is_final:
+            final_ipa_types.append(ipa_type)
+            if ipa_type.is_ipa:
+                stub_file = stub_out / cls.dir / f"{cls.name}Constructor.qll"
+                if not stub_file.is_file() or _is_generated_stub(stub_file):
+                    renderer.render(ql.Ipa.ConstructorStub(ipa_type), stub_file)
+                constructor_imports.append(get_import(stub_file, opts.swift_dir))
+        else:
+            non_final_ipa_types.append(ipa_type)
+
+    renderer.render(ql.Ipa.Types(schema.root_class_name, final_ipa_types, non_final_ipa_types), out / "IpaTypes.qll")
+    renderer.render(ql.ImportList(constructor_imports), out / "IpaConstructors.qll")
+
     renderer.cleanup(existing)
     if opts.ql_format:
         format(opts.codeql_binary, renderer.written)

swift/codegen/lib/ql.py

@@ -14,7 +14,7 @@
 
 import pathlib
 from dataclasses import dataclass, field
-from typing import List, ClassVar
+from typing import List, ClassVar, Union
 
 import inflection
 
@@ -151,3 +151,93 @@ class PropertyTester:
 @dataclass
 class MissingTestInstructions:
     template: ClassVar = 'ql_test_missing'
+
+
+class Ipa:
+    @dataclass
+    class Class:
+        is_final: ClassVar = False
+
+        name: str
+        first: bool = False
+
+    @dataclass
+    class FinalClass(Class):
+        is_final: ClassVar = True
+        is_ipa_from: ClassVar = False
+        is_ipa_on: ClassVar = False
+        is_db: ClassVar = False
+
+        @property
+        def is_ipa(self):
+            return self.is_ipa_on or self.is_ipa_from
+
+    @dataclass
+    class Param:
+        param: str
+        type: str
+        first: bool = False
+
+    @dataclass
+    class FinalClassIpaFrom(FinalClass):
+        is_ipa_from: ClassVar = True
+
+        type: str = None
+
+        @property
+        def params(self):
+            return [Ipa.Param("id", self.type, first=True)]
+
+    @dataclass
+    class FinalClassIpaOn(FinalClass):
+        is_ipa_on: ClassVar = True
+
+        params: List["Ipa.Param"] = field(default_factory=list)
+
+        def __post_init__(self):
+            if self.params:
+                self.params[0].first = True
+
+    @dataclass
+    class FinalClassDb(FinalClass):
+        is_db: ClassVar = True
+
+        subtracted_ipa_types: List["Ipa.Class"] = field(default_factory=list)
+
+        def subtract_type(self, type: str):
+            self.subtracted_ipa_types.append(Ipa.Class(type, first=not self.subtracted_ipa_types))
+
+        @property
+        def has_subtracted_ipa_types(self):
+            return bool(self.subtracted_ipa_types)
+
+        @property
+        def db_id(self):
+            return "@" + inflection.underscore(self.name)
+
+    @dataclass
+    class NonFinalClass(Class):
+        derived: List["Ipa.Class"] = field(default_factory=list)
+
+        def __post_init__(self):
+            self.derived = [Ipa.Class(c) for c in self.derived]
+            if self.derived:
+                self.derived[0].first = True
+
+    @dataclass
+    class Types:
+        template: ClassVar = "ql_ipa_types"
+
+        root: str
+        final_classes: List["Ipa.FinalClass"] = field(default_factory=list)
+        non_final_classes: List["Ipa.NonFinalClass"] = field(default_factory=list)
+
+        def __post_init__(self):
+            if self.final_classes:
+                self.final_classes[0].first = True
+
+    @dataclass
+    class ConstructorStub:
+        template: ClassVar = "ql_ipa_constructor_stub"
+
+        cls: Union["Ipa.FinalClassIpaFrom", "Ipa.FinalClassIpaOn"]

swift/codegen/lib/schema.py

@@ -2,7 +2,6 @@
 
 import pathlib
 import re
-import typing
 from dataclasses import dataclass, field
 from typing import List, Set, Union, Dict, ClassVar
 
@@ -57,6 +56,12 @@ class PredicateProperty(Property):
     is_predicate: ClassVar = True
 
 
+@dataclass
+class IpaInfo:
+    from_class: str = None
+    on_arguments: Dict[str, str] = None
+
+
 @dataclass
 class Class:
     name: str
@@ -65,6 +70,7 @@ class Class:
     properties: List[Property] = field(default_factory=list)
     dir: pathlib.Path = pathlib.Path()
     pragmas: List[str] = field(default_factory=list)
+    ipa: IpaInfo = None
 
 
 @dataclass
@@ -107,6 +113,11 @@ def _parse_property(name: str, data: Union[str, Dict[str, _StrOrList]], is_child
         return SingleProperty(name, type, is_child=is_child, pragmas=pragmas)
 
 
+def _parse_ipa(data: Dict[str, Union[str, Dict[str, str]]]):
+    return IpaInfo(from_class=data.get("from"),
+                   on_arguments=data.get(True))  # 'on' is parsed as boolean True in yaml
+
+
 class _DirSelector:
     """ Default output subdirectory selector for generated QL files, based on the `_directories` global field"""
 
@@ -145,6 +156,8 @@ def load(path):
                 cls.properties.extend(_parse_property(kk, vv, is_child=True) for kk, vv in v.items())
             elif k == "_pragma":
                 cls.pragmas = _auto_list(v)
+            elif k == "_ipa":
+                cls.ipa = _parse_ipa(v)
             else:
                 raise Error(f"unknown metadata {k} for class {name}")
         if not cls.bases and cls.name != root_class_name:

swift/codegen/schema.yml

@@ -792,6 +792,12 @@ OptionalTryExpr:
 TryExpr:
   _extends: AnyTryExpr
 
+MethodCallExpr:
+  _extends: ApplyExpr
+  _ipa:
+    from: CallExpr
+  qualifier: Expr
+
 BinaryExpr:
   _extends: ApplyExpr
 

swift/codegen/templates/ql_class.mustache

@@ -1,5 +1,4 @@
 // generated by {{generator}}
-
 {{#imports}}
 import {{.}}
 {{/imports}}

swift/codegen/templates/ql_ipa_constructor_stub.mustache

@@ -0,0 +1,2 @@
+// generated by {{generator}}, remove this comment if you wish to edit this file
+predicate construct{{cls.name}}({{#cls.params}}{{^first}}, {{/first}}{{type}} {{param}}{{/cls.params}}) { none() }

swift/codegen/templates/ql_ipa_types.mustache

@@ -0,0 +1,23 @@
+private import codeql.swift.elements.IpaConstructors
+
+cached module Cached {
+  cached newtype T{{root}} =
+    {{#final_classes}}
+    {{^first}}
+    or
+    {{/first}}
+    {{#is_ipa_from}}
+    T{{name}}({{type}} id) { construct{{name}}(id) }
+    {{/is_ipa_from}}
+    {{#is_ipa_on}}
+    T{{name}}({{#params}}{{^first}}, {{/first}}{{type}} {{param}}{{/params}}) { construct{{name}}({{#params}}{{^first}}, {{/first}}{{param}}{{/params}}) }
+    {{/is_ipa_on}}
+    {{#is_db}}
+    T{{name}}({{db_id}} id){{#has_subtracted_ipa_types}} { {{#subtracted_ipa_types}}{{^first}} and {{/first}}not construct{{name}}(id){{/subtracted_ipa_types}} }{{/has_subtracted_ipa_types}}
+    {{/is_db}}
+    {{/final_classes}}
+
+  {{#non_final_classes}}
+  class T{{name}} = {{#derived}}{{^first}} or {{/first}}T{{name}}{{/derived}};
+  {{/non_final_classes}}
+}

swift/codegen/templates/ql_stub.mustache

@@ -1,5 +1,4 @@
 // generated by {{generator}}, remove this comment if you wish to edit this file
-
 private import {{base_import}}
 
 class {{name}} extends {{name}}Base {}

swift/codegen/test/test_qlgen.py

@@ -101,6 +101,7 @@ def _filter_generated_classes(ret, output_test_files=False):
             str(f): ret[ql_test_output_path() / f]
             for f in test_files
         }
+    base_files -= {pathlib.Path("IpaTypes.qll"), pathlib.Path("IpaConstructors.qll")}
     assert stub_files == base_files
     return {
         str(f): (ret[stub_path() / f], ret[ql_output_path() / f])
@@ -128,6 +129,8 @@ def test_empty(generate):
     assert generate([]) == {
         import_file(): ql.ImportList(),
         children_file(): ql.GetParentImplementation(),
+        ql_output_path() / "IpaTypes.qll": ql.Ipa.Types(schema.root_class_name),
+        ql_output_path() / "IpaConstructors.qll": ql.ImportList(),
     }
 
 

swift/codegen/test/test_schema.py

@@ -303,5 +303,34 @@ def test_class_with_unknown_metadata(load):
 """)
 
 
+def test_ipa_class_from(load):
+    ret = load("""
+MyClass:
+    _ipa:
+        from: A
+""")
+    assert ret.classes == [
+        schema.Class(root_name, derived={'MyClass'}),
+        schema.Class('MyClass', bases={root_name}, ipa=schema.IpaInfo(from_class="A")),
+    ]
+
+
+def test_ipa_class_on(load):
+    ret = load("""
+MyClass:
+    _ipa:
+        on: 
+            x: A
+            y: int
+""")
+    assert ret.classes == [
+        schema.Class(root_name, derived={'MyClass'}),
+        schema.Class('MyClass', bases={root_name}, ipa=schema.IpaInfo(on_arguments={"x": "A", "y": "int"})),
+    ]
+
+
+# TODO rejection tests and implementation for malformed `_ipa` clauses
+
+
 if __name__ == '__main__':
     sys.exit(pytest.main([__file__] + sys.argv[1:]))