Merge pull request #10211 from github/redsun82/swift-generate-printast-first-step

Swift: generate indexed parent/child relationship

Author: MathiasVP

swift/codegen/generators/qlgen.py

@@ -29,11 +29,15 @@ class ModifiedStubMarkedAsGeneratedError(Error):
     pass
 
 
-def get_ql_property(cls: schema.Class, prop: schema.Property) -> ql.Property:
+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",
         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 +73,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,
@@ -116,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()))
 
 
@@ -227,10 +238,15 @@ 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()}
-    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 +267,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):

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
@@ -35,7 +36,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,13 +66,26 @@ 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 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()
@@ -82,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
 
@@ -95,9 +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) -> 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

@@ -2,33 +2,88 @@
 
 import codeql.swift.elements
 
-/**
- * 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.
- */
+private module Impl {
+{{#classes}}
+    private 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 offset 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}}
+            {{! 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
+            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}}
+        ))
+    {{/root}}
+    }
+
+{{/classes}}
 cached
-Element getAnImmediateChild(Element e) {
-  // 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}}
-  {{/classes}}
+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 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)
+}
+
+/**
+* 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)
 }
 
 /**
- * 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)
+* 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)
 }

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
 
@@ -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:]))