Swift: add indexes to generated parent-child library

redsun82 · redsun82 · commit b5aad04f6593 · 2022-08-29T17:10:46.000+02:00
diff --git a/swift/codegen/generators/qlgen.py b/swift/codegen/generators/qlgen.py
@@ -29,11 +29,11 @@ class ModifiedStubMarkedAsGeneratedError(Error):
     pass
 
 
-def get_ql_property(cls: schema.Class, prop: schema.Property) -> ql.Property:
+def get_ql_property(cls: schema.Class, prop: schema.Property, prev_child: str = "") -> ql.Property:
     args = dict(
         type=prop.type if not prop.is_predicate else "predicate",
         qltest_skip="qltest_skip" in prop.pragmas,
-        is_child=prop.is_child,
+        prev_child=prev_child if prop.is_child else None,
         is_optional=prop.is_optional,
         is_predicate=prop.is_predicate,
     )
@@ -69,11 +69,18 @@ def get_ql_property(cls: schema.Class, prop: schema.Property) -> ql.Property:
 
 def get_ql_class(cls: schema.Class, lookup: typing.Dict[str, schema.Class]):
     pragmas = {k: True for k in cls.pragmas if k.startswith("ql")}
+    prev_child = ""
+    properties = []
+    for p in cls.properties:
+        prop = get_ql_property(cls, p, prev_child)
+        if prop.is_child:
+            prev_child = prop.singular
+        properties.append(prop)
     return ql.Class(
         name=cls.name,
         bases=cls.bases,
         final=not cls.derived,
-        properties=[get_ql_property(cls, p) for p in cls.properties],
+        properties=properties,
         dir=cls.dir,
         ipa=bool(cls.ipa),
         **pragmas,
@@ -230,7 +237,8 @@ def generate(opts, renderer):
     imports = {}
 
     inheritance_graph = {name: cls.bases for name, cls in data.classes.items()}
-    db_classes = [classes[name] for name in toposort_flatten(inheritance_graph) if not classes[name].ipa]
+    toposorted_names = toposort_flatten(inheritance_graph)
+    db_classes = [classes[name] for name in toposorted_names if not classes[name].ipa]
     renderer.render(ql.DbClasses(db_classes), out / "Raw.qll")
 
     classes_by_dir_and_name = sorted(classes.values(), key=lambda cls: (cls.dir, cls.name))
@@ -251,7 +259,8 @@ def generate(opts, renderer):
     include_file = stub_out.with_suffix(".qll")
     renderer.render(ql.ImportList(list(imports.values())), include_file)
 
-    renderer.render(ql.GetParentImplementation(classes_by_dir_and_name), out / 'GetImmediateParent.qll')
+    toposorted_classes = [classes[name] for name in toposorted_names]
+    renderer.render(ql.GetParentImplementation(toposorted_classes), out / 'ParentChild.qll')
 
     for c in data.classes.values():
         if _should_skip_qltest(c, data.classes):
diff --git a/swift/codegen/lib/ql.py b/swift/codegen/lib/ql.py
@@ -35,7 +35,7 @@ class Property:
     first: bool = False
     is_optional: bool = False
     is_predicate: bool = False
-    is_child: bool = False
+    prev_child: Optional[str] = None
     qltest_skip: bool = False
 
     def __post_init__(self):
@@ -65,6 +65,10 @@ def is_repeated(self):
     def is_single(self):
         return not (self.is_optional or self.is_repeated or self.is_predicate)
 
+    @property
+    def is_child(self):
+        return self.prev_child is not None
+
 
 @dataclass
 class Class:
@@ -98,6 +102,10 @@ def path(self) -> pathlib.Path:
     def db_id(self):
         return "@" + inflection.underscore(self.name)
 
+    @property
+    def has_children(self):
+        return any(p.is_child for p in self.properties)
+
 
 @dataclass
 class Stub:
diff --git a/swift/codegen/templates/ql_parent.mustache b/swift/codegen/templates/ql_parent.mustache
@@ -2,33 +2,82 @@
 
 import codeql.swift.elements
 
+module Impl {
+{{#classes}}
+int getMaximumChildrenIndex{{name}}({{name}} e) {
+  {{#root}}e = e and{{/root}}
+  result = 0
+    {{#bases}}
+    + getMaximumChildrenIndex{{.}}(e)
+    {{/bases}}
+    {{#properties}}
+    {{#is_child}}
+    + 1{{#is_repeated}}+ max(int i | exists(e.getImmediate{{singular}}(i)) | i){{/is_repeated}}
+    {{/is_child}}
+    {{/properties}}
+}
+
+{{/classes}}
 /**
  * Gets any of the "immediate" children of `e`. "Immediate" means not taking into account node resolution: for example
  * if the AST child is the first of a series of conversions that would normally be hidden away, this will select the
  * next conversion down the hidden AST tree instead of the corresponding fully uncoverted node at the bottom.
  * Outside this module this file is mainly intended to be used to test uniqueness of parents.
  */
 cached
-Element getAnImmediateChild(Element e) {
+Element getImmediateChild(Element e, int index, string partialAccessor) {
   // why does this look more complicated than it should?
   // * none() simplifies generation, as we can append `or ...` without a special case for the first item
   none()
   {{#classes}}
-  {{#properties}}
-  {{#is_child}}
-  or
-  result = e.({{name}}).getImmediate{{singular}}({{#is_repeated}}_{{/is_repeated}})
-  {{/is_child}}
-  {{/properties}}
+      {{#has_children}}
+      or
+      exists(int n{{#properties}}{{#is_child}}, int n{{singular}}{{/is_child}}{{/properties}} |
+        n = 0{{#bases}} + getMaximumChildrenIndex{{.}}(e){{/bases}}
+        {{#properties}}
+        {{#is_child}}
+        and n{{singular}} = n{{prev_child}} + 1{{#is_repeated}} + max(int i | i = 0 or exists(e.({{name}}).getImmediate{{singular}}(i)) | i){{/is_repeated}}
+        {{/is_child}}
+        {{/properties}}
+        and (
+            none()
+            {{#properties}}
+            {{#is_child}}
+            or
+            {{#is_repeated}}
+                result = e.({{name}}).getImmediate{{singular}}(index - n{{prev_child}}) and partialAccessor = "{{singular}}(" + (index - n{{prev_child}}).toString() + ")"
+            {{/is_repeated}}
+            {{^is_repeated}}
+                index = n{{prev_child}} and result = e.({{name}}).getImmediate{{singular}}() and partialAccessor = "{{singular}}()"
+            {{/is_repeated}}
+            {{/is_child}}
+            {{/properties}}
+      ))
+      {{/has_children}}
   {{/classes}}
 }
+}
 
 /**
  * Gets the "immediate" parent of `e`. "Immediate" means not taking into account node resolution: for example
  * if `e` has conversions, `getImmediateParent(e)` will give the bottom conversion in the hidden AST.
  */
- Element getImmediateParent(Element e) {
-   // `unique` is used here to tell the optimizer that there is in fact only one result
-   // this is tested by the `library-tests/parent/no_double_parents.ql` test
-   result = unique(Element x | e = getAnImmediateChild(x) | x)
+Element getImmediateParent(Element e) {
+    // `unique` is used here to tell the optimizer that there is in fact only one result
+    // this is tested by the `library-tests/parent/no_double_parents.ql` test
+    result = unique(Element x | e = Impl::getImmediateChild(x, _, _) | x)
+}
+
+/**
+* Gets the immediate child indexed at `index`. Indexes are not guaranteed to be contiguous, but are guaranteed to be distinct. `accessor` is bound the the method giving the given child.
+*/
+Element getImmediateChildAndAccessor(Element e, int index, string accessor) {
+    exists(string partialAccessor | result = Impl::getImmediateChild(e, index, partialAccessor) and accessor = "getImmediate" + partialAccessor)
+}
+
+/**
+ * Gets the child indexed at `index`. Indexes are not guaranteed to be contiguous, but are guaranteed to be distinct. `accessor` is bound the the method giving the given child. Node resolution is carried out.
+ */
+Element getChildAndAccessor(Element e, int index, string accessor) {
+    exists(string partialAccessor | result = Impl::getImmediateChild(e, index, partialAccessor).resolve() and accessor = "get" + partialAccessor)
 }
diff --git a/swift/codegen/test/test_ql.py b/swift/codegen/test/test_ql.py
@@ -105,5 +105,27 @@ def test_non_root_class():
     assert not cls.root
 
 
+@pytest.mark.parametrize("prev_child,is_child", [(None, False), ("", True), ("x", True)])
+def test_is_child(prev_child, is_child):
+    p = ql.Property("Foo", "int", prev_child=prev_child)
+    assert p.is_child is is_child
+
+
+def test_empty_class_no_children():
+    cls = ql.Class("Class", properties=[])
+    assert cls.has_children is False
+
+
+def test_class_no_children():
+    cls = ql.Class("Class", properties=[ql.Property("Foo", "int"), ql.Property("Bar", "string")])
+    assert cls.has_children is False
+
+
+def test_class_with_children():
+    cls = ql.Class("Class", properties=[ql.Property("Foo", "int"), ql.Property("Child", "x", prev_child=""),
+                                        ql.Property("Bar", "string")])
+    assert cls.has_children is True
+
+
 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
@@ -29,7 +29,7 @@ def ql_test_output_path(): return paths.swift_dir / "ql/test/path"
 def import_file(): return stub_path().with_suffix(".qll")
 
 
-def children_file(): return ql_output_path() / "GetImmediateParent.qll"
+def children_file(): return ql_output_path() / "ParentChild.qll"
 
 
 stub_import_prefix = "stub.path."
@@ -207,6 +207,42 @@ def test_single_property(generate_classes):
     }
 
 
+def test_children(generate_classes):
+    assert generate_classes([
+        schema.Class("MyObject", properties=[
+            schema.SingleProperty("a", "int"),
+            schema.SingleProperty("child1", "int", is_child=True),
+            schema.RepeatedProperty("b", "int"),
+            schema.RepeatedProperty("child2", "int", is_child=True),
+            schema.OptionalProperty("c", "int"),
+            schema.OptionalProperty("child3", "int", is_child=True),
+            schema.RepeatedOptionalProperty("d", "int"),
+            schema.RepeatedOptionalProperty("child4", "int", is_child=True),
+        ]),
+    ]) == {
+        "MyObject.qll": (ql.Stub(name="MyObject", base_import=gen_import_prefix + "MyObject"),
+                         ql.Class(name="MyObject", final=True,
+                                  properties=[
+                                      ql.Property(singular="A", type="int", tablename="my_objects",
+                                                  tableparams=["this", "result", "_"]),
+                                      ql.Property(singular="Child1", type="int", tablename="my_objects",
+                                                  tableparams=["this", "_", "result"], prev_child=""),
+                                      ql.Property(singular="B", plural="Bs", type="int", tablename="my_object_bs",
+                                                  tableparams=["this", "index", "result"]),
+                                      ql.Property(singular="Child2", plural="Child2s", type="int", tablename="my_object_child2s",
+                                                  tableparams=["this", "index", "result"], prev_child="Child1"),
+                                      ql.Property(singular="C", type="int", tablename="my_object_cs",
+                                                  tableparams=["this", "result"], is_optional=True),
+                                      ql.Property(singular="Child3", type="int", tablename="my_object_child3s",
+                                                  tableparams=["this", "result"], is_optional=True, prev_child="Child2"),
+                                      ql.Property(singular="D", plural="Ds", type="int", tablename="my_object_ds",
+                                                  tableparams=["this", "index", "result"], is_optional=True),
+                                      ql.Property(singular="Child4", plural="Child4s", type="int", tablename="my_object_child4s",
+                                                  tableparams=["this", "index", "result"], is_optional=True, prev_child="Child3"),
+                                  ])),
+    }
+
+
 def test_single_properties(generate_classes):
     assert generate_classes([
         schema.Class("MyObject", properties=[
@@ -228,8 +264,8 @@ def test_single_properties(generate_classes):
     }
 
 
-@pytest.mark.parametrize("is_child", [False, True])
-def test_optional_property(generate_classes, is_child):
+@pytest.mark.parametrize("is_child,prev_child", [(False, None), (True, "")])
+def test_optional_property(generate_classes, is_child, prev_child):
     assert generate_classes([
         schema.Class("MyObject", properties=[
             schema.OptionalProperty("foo", "bar", is_child=is_child)]),
@@ -238,27 +274,27 @@ def test_optional_property(generate_classes, is_child):
                          ql.Class(name="MyObject", final=True, properties=[
                              ql.Property(singular="Foo", type="bar", tablename="my_object_foos",
                                          tableparams=["this", "result"],
-                                         is_optional=True, is_child=is_child),
+                                         is_optional=True, prev_child=prev_child),
                          ])),
     }
 
 
-@pytest.mark.parametrize("is_child", [False, True])
-def test_repeated_property(generate_classes, is_child):
+@pytest.mark.parametrize("is_child,prev_child", [(False, None), (True, "")])
+def test_repeated_property(generate_classes, is_child, prev_child):
     assert generate_classes([
         schema.Class("MyObject", properties=[
             schema.RepeatedProperty("foo", "bar", is_child=is_child)]),
     ]) == {
         "MyObject.qll": (ql.Stub(name="MyObject", base_import=gen_import_prefix + "MyObject"),
                          ql.Class(name="MyObject", final=True, properties=[
                              ql.Property(singular="Foo", plural="Foos", type="bar", tablename="my_object_foos",
-                                         tableparams=["this", "index", "result"], is_child=is_child),
+                                         tableparams=["this", "index", "result"], prev_child=prev_child),
                          ])),
     }
 
 
-@pytest.mark.parametrize("is_child", [False, True])
-def test_repeated_optional_property(generate_classes, is_child):
+@pytest.mark.parametrize("is_child,prev_child", [(False, None), (True, "")])
+def test_repeated_optional_property(generate_classes, is_child, prev_child):
     assert generate_classes([
         schema.Class("MyObject", properties=[
             schema.RepeatedOptionalProperty("foo", "bar", is_child=is_child)]),
@@ -267,7 +303,7 @@ def test_repeated_optional_property(generate_classes, is_child):
                          ql.Class(name="MyObject", final=True, properties=[
                              ql.Property(singular="Foo", plural="Foos", type="bar", tablename="my_object_foos",
                                          tableparams=["this", "index", "result"], is_optional=True,
-                                         is_child=is_child),
+                                         prev_child=prev_child),
                          ])),
     }
 
@@ -286,8 +322,8 @@ def test_predicate_property(generate_classes):
     }
 
 
-@pytest.mark.parametrize("is_child", [False, True])
-def test_single_class_property(generate_classes, is_child):
+@pytest.mark.parametrize("is_child,prev_child", [(False, None), (True, "")])
+def test_single_class_property(generate_classes, is_child, prev_child):
     assert generate_classes([
         schema.Class("MyObject", properties=[
             schema.SingleProperty("foo", "Bar", is_child=is_child)]),
@@ -299,7 +335,7 @@ def test_single_class_property(generate_classes, is_child):
                 ql.Property(singular="Foo", type="Bar", tablename="my_objects",
                             tableparams=[
                                 "this", "result"],
-                            is_child=is_child),
+                            prev_child=prev_child),
             ],
         )),
         "Bar.qll": (ql.Stub(name="Bar", base_import=gen_import_prefix + "Bar"),
diff --git a/swift/ql/lib/codeql/swift/generated/GetImmediateParent.qll b/swift/ql/lib/codeql/swift/generated/GetImmediateParent.qll
diff --git a/swift/ql/lib/codeql/swift/generated/ParentChild.qll b/swift/ql/lib/codeql/swift/generated/ParentChild.qll
