Swift: fix ParentChild generation

There was an issue in case multiple inheritance from classes with
children was involved, where indexes would overlap.

The generated code structure has been reshuffled a bit, with
`Impl::getImmediateChildOf<Class>` predicates giving 0-based children
for a given class, including those coming from bases, and the final
`Impl::getImmediateChild` disjuncting the above on final classes only.

This removes the need of `getMaximumChildrenIndex<Class>`, and also
removes the code scanning alerts.

Also, comments were fixed addressing the review.

Author: redsun82

swift/codegen/generators/qlgen.py

@@ -29,6 +29,10 @@ class ModifiedStubMarkedAsGeneratedError(Error):
     pass
 
 
+class RootElementHasChildren(Error):
+    pass
+
+
 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",
@@ -123,14 +127,14 @@ def get_import(file: pathlib.Path, swift_dir: pathlib.Path):
     return str(stem).replace("/", ".")
 
 
-def get_types_used_by(cls: ql.Class):
+def get_types_used_by(cls: ql.Class) -> typing.Iterable[str]:
     for b in cls.bases:
-        yield b
+        yield b.base
     for p in cls.properties:
         yield p.type
 
 
-def get_classes_used_by(cls: ql.Class):
+def get_classes_used_by(cls: ql.Class) -> typing.List[str]:
     return sorted(set(t for t in get_types_used_by(cls) if t[0].isupper()))
 
 
@@ -234,6 +238,10 @@ def generate(opts, renderer):
     data = schema.load(input)
 
     classes = {name: get_ql_class(cls, data.classes) for name, cls in data.classes.items()}
+    # element root is absent in tests
+    if schema.root_class_name in classes and classes[schema.root_class_name].has_children:
+        raise RootElementHasChildren
+
     imports = {}
 
     inheritance_graph = {name: cls.bases for name, cls in data.classes.items()}

swift/codegen/lib/ql.py

@@ -14,6 +14,7 @@
 
 import pathlib
 from dataclasses import dataclass, field
+import itertools
 from typing import List, ClassVar, Union, Optional
 
 import inflection
@@ -70,12 +71,21 @@ def is_child(self):
         return self.prev_child is not None
 
 
+@dataclass
+class Base:
+    base: str
+    prev: str = ""
+
+    def __str__(self):
+        return self.base
+
+
 @dataclass
 class Class:
     template: ClassVar = 'ql_class'
 
     name: str
-    bases: List[str] = field(default_factory=list)
+    bases: List[Base] = field(default_factory=list)
     final: bool = False
     properties: List[Property] = field(default_factory=list)
     dir: pathlib.Path = pathlib.Path()
@@ -86,7 +96,8 @@ class Class:
     ipa: bool = False
 
     def __post_init__(self):
-        self.bases = sorted(self.bases)
+        bases = sorted(str(b) for b in self.bases)
+        self.bases = [Base(str(b), prev) for b, prev in zip(bases, itertools.chain([""], bases))]
         if self.properties:
             self.properties[0].first = True
 
@@ -99,13 +110,17 @@ def path(self) -> pathlib.Path:
         return self.dir / self.name
 
     @property
-    def db_id(self):
+    def db_id(self) -> str:
         return "@" + inflection.underscore(self.name)
 
     @property
-    def has_children(self):
+    def has_children(self) -> bool:
         return any(p.is_child for p in self.properties)
 
+    @property
+    def last_base(self) -> str:
+        return self.bases[-1].base if self.bases else ""
+
 
 @dataclass
 class Stub:

swift/codegen/templates/ql_parent.mustache

@@ -4,80 +4,86 @@ import codeql.swift.elements
 
 private 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 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}}
-      {{#has_children}}
-      or
-      exists(int n{{#properties}}{{#is_child}}, int n{{singular}}{{/is_child}}{{/properties}} |
-        n = 0{{#bases}} + getMaximumChildrenIndex{{.}}(e){{/bases}}
+    Element getImmediateChildOf{{name}}({{name}} e, int index, string partialPredicateCall) {
+    {{! avoid unused argument warnings on root element, assuming the root element has no children }}
+    {{#root}}none(){{/root}}
+    {{^root}}
+    {{! b is the base offset 0, for ease of generation }}
+    {{! b<base> is constructed to be strictly greater than the indexes required for children coming from <base> }}
+    {{! n is the base offset for direct children, equal to the last base offsets from above }}
+    {{! n<child> is constructed to be strictly greater than the indexes for <child> children }}
+        exists(int b{{#bases}}, int b{{.}}{{/bases}}, int n{{#properties}}{{#is_child}}, int n{{singular}}{{/is_child}}{{/properties}} |
+        b = 0
+        {{#bases}}
+            and
+            b{{.}} = b{{prev}} + 1 + max(int i | i = -1 or exists(getImmediateChildOf{{.}}(e, i, _)) | i)
+        {{/bases}}
+        and
+        n = b{{last_base}}
         {{#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}}
+            {{! n<child> is defined on top of the previous definition }}
+            {{! for single and optional properties it adds 1 (regardless of whether the optional property exists) }}
+            {{! for repeated it adds 1 + the maximum index (which works for repeated optional as well) }}
+                and
+                n{{singular}} = n{{prev_child}} + 1{{#is_repeated}}+ max(int i | i = -1 or exists(e.getImmediate{{singular}}(i)) | i){{/is_repeated}}
+            {{/is_child}}
+        {{/properties}} and (
+        none()
+        {{#bases}}
             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}}
+            result = getImmediateChildOf{{.}}(e, index - b{{prev}}, partialPredicateCall)
+        {{/bases}}
+        {{#properties}}
+            {{#is_child}}
+                or
+                {{#is_repeated}}
+                    result = e.getImmediate{{singular}}(index - n{{prev_child}}) and partialPredicateCall = "{{singular}}(" + (index - n{{prev_child}}).toString() + ")"
+                {{/is_repeated}}
+                {{^is_repeated}}
+                    index = n{{prev_child}} and result = e.getImmediate{{singular}}() and partialPredicateCall = "{{singular}}()"
+                {{/is_repeated}}
             {{/is_child}}
-            {{/properties}}
-      ))
-      {{/has_children}}
-  {{/classes}}
+        {{/properties}}
+        ))
+    {{/root}}
+    }
+
+{{/classes}}
+cached
+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}}
+    {{#final}}
+        or
+        result = getImmediateChildOf{{name}}(e, index, partialAccessor)
+    {{/final}}
+{{/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.
- */
+* 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 innermost 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 = Impl::getImmediateChild(x, _, _) | x)
+// `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.
+* Gets the immediate child indexed at `index`. Indexes are not guaranteed to be contiguous, but are guaranteed to be distinct. `accessor` is bound the member predicate call resulting in the given child.
 */
 Element getImmediateChildAndAccessor(Element e, int index, string accessor) {
-    exists(string partialAccessor | result = Impl::getImmediateChild(e, index, partialAccessor) and accessor = "getImmediate" + partialAccessor)
+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.
- */
+* Gets the child indexed at `index`. Indexes are not guaranteed to be contiguous, but are guaranteed to be distinct. `accessor` is bound the member predicate call resulting in the given child.
+*/
 Element getChildAndAccessor(Element e, int index, string accessor) {
-    exists(string partialAccessor | result = Impl::getImmediateChild(e, index, partialAccessor).resolve() and accessor = "get" + partialAccessor)
+exists(string partialAccessor | result = Impl::getImmediateChild(e, index, partialAccessor).resolve() and accessor = "get" + partialAccessor)
 }

swift/codegen/test/test_ql.py

@@ -78,9 +78,9 @@ def test_property_predicate_getter():
     assert prop.getter == "prop"
 
 
-def test_class_sorts_bases():
+def test_class_processes_bases():
     bases = ["B", "Ab", "C", "Aa"]
-    expected = ["Aa", "Ab", "B", "C"]
+    expected = [ql.Base("Aa"), ql.Base("Ab", prev="Aa"), ql.Base("B", prev="Ab"), ql.Base("C", prev="B")]
     cls = ql.Class("Foo", bases=bases)
     assert cls.bases == expected
 

swift/codegen/test/test_qlgen.py

@@ -357,6 +357,13 @@ def test_class_dir(generate_classes):
     }
 
 
+def test_root_element_cannot_have_children(generate_classes):
+    with pytest.raises(qlgen.RootElementHasChildren):
+        generate_classes([
+            schema.Class(schema.root_class_name, properties=[schema.SingleProperty("x", is_child=True)])
+        ])
+
+
 def test_class_dir_imports(generate_import_list):
     dir = pathlib.Path("another/rel/path")
     assert generate_import_list([