Swift: extract all ReferenceStorageTypes

Author: redsun82

swift/codegen/schema.yml

@@ -121,6 +121,7 @@ ExistentialType:
 
 ReferenceStorageType:
   _extends: Type
+  referent_type: Type
 
 SilBlockStorageType:
   _extends: Type

swift/extractor/visitors/TypeVisitor.cpp

@@ -278,4 +278,30 @@ codeql::InOutType TypeVisitor::translateInOutType(const swift::InOutType& type)
   return entry;
 }
 
+codeql::UnmanagedStorageType TypeVisitor::translateUnmanagedStorageType(
+    const swift::UnmanagedStorageType& type) {
+  codeql::UnmanagedStorageType entry{dispatcher_.assignNewLabel(type)};
+  fillReferenceStorageType(type, entry);
+  return entry;
+}
+
+codeql::UnownedStorageType TypeVisitor::translateUnownedStorageType(
+    const swift::UnownedStorageType& type) {
+  codeql::UnownedStorageType entry{dispatcher_.assignNewLabel(type)};
+  fillReferenceStorageType(type, entry);
+  return entry;
+}
+
+codeql::WeakStorageType TypeVisitor::translateWeakStorageType(const swift::WeakStorageType& type) {
+  codeql::WeakStorageType entry{dispatcher_.assignNewLabel(type)};
+  fillReferenceStorageType(type, entry);
+  return entry;
+}
+
+void TypeVisitor::fillReferenceStorageType(const swift::ReferenceStorageType& type,
+                                           codeql::ReferenceStorageType& entry) {
+  entry.referent_type = dispatcher_.fetchLabel(type.getReferentType());
+  fillType(type, entry);
+}
+
 }  // namespace codeql

swift/extractor/visitors/TypeVisitor.h

@@ -43,12 +43,18 @@ class TypeVisitor : public TypeVisitorBase<TypeVisitor> {
   codeql::VariadicSequenceType translateVariadicSequenceType(
       const swift::VariadicSequenceType& type);
   codeql::InOutType translateInOutType(const swift::InOutType& type);
+  codeql::UnmanagedStorageType translateUnmanagedStorageType(
+      const swift::UnmanagedStorageType& type);
+  codeql::WeakStorageType translateWeakStorageType(const swift::WeakStorageType& type);
+  codeql::UnownedStorageType translateUnownedStorageType(const swift::UnownedStorageType& type);
 
  private:
   void fillType(const swift::TypeBase& type, codeql::Type& entry);
   void fillArchetypeType(const swift::ArchetypeType& type, codeql::ArchetypeType& entry);
   void fillUnarySyntaxSugarType(const swift::UnarySyntaxSugarType& type,
                                 codeql::UnarySyntaxSugarType& entry);
+  void fillReferenceStorageType(const swift::ReferenceStorageType& type,
+                                codeql::ReferenceStorageType& entry);
   void emitAnyFunctionType(const swift::AnyFunctionType* type, TrapLabel<AnyFunctionTypeTag> label);
   void emitBoundGenericType(swift::BoundGenericType* type, TrapLabel<BoundGenericTypeTag> label);
   void emitAnyGenericType(swift::AnyGenericType* type, TrapLabel<AnyGenericTypeTag> label);

swift/ql/lib/codeql/swift/generated/type/ReferenceStorageType.qll

@@ -1,4 +1,11 @@
 // generated by codegen/codegen.py
 import codeql.swift.elements.type.Type
 
-class ReferenceStorageTypeBase extends @reference_storage_type, Type { }
+class ReferenceStorageTypeBase extends @reference_storage_type, Type {
+  Type getReferentType() {
+    exists(Type x |
+      reference_storage_types(this, x) and
+      result = x.resolve()
+    )
+  }
+}

swift/ql/lib/swift.dbscheme

@@ -285,6 +285,12 @@ existential_types(  //dir=type
 | @weak_storage_type
 ;
 
+#keyset[id]
+reference_storage_types(  //dir=type
+  int id: @reference_storage_type ref,
+  int referent_type: @type ref
+);
+
 sil_block_storage_types(  //dir=type
   unique int id: @sil_block_storage_type
 );

swift/ql/test/extractor-tests/generated/type/UnmanagedStorageType/MISSING_SOURCE.txt

@@ -0,0 +1 @@
+| A? | getDiagnosticsName: | A? | getCanonicalType: | Optional<A> | getReferentType: | A? |

swift/ql/test/extractor-tests/generated/type/UnmanagedStorageType/UnmanagedStorageType.expected

@@ -0,0 +1,13 @@
+// generated by codegen/codegen.py
+import codeql.swift.elements
+import TestUtils
+
+from UnmanagedStorageType x, string getDiagnosticsName, Type getCanonicalType, Type getReferentType
+where
+  toBeTested(x) and
+  not x.isUnknown() and
+  getDiagnosticsName = x.getDiagnosticsName() and
+  getCanonicalType = x.getCanonicalType() and
+  getReferentType = x.getReferentType()
+select x, "getDiagnosticsName:", getDiagnosticsName, "getCanonicalType:", getCanonicalType,
+  "getReferentType:", getReferentType

swift/ql/test/extractor-tests/generated/type/UnmanagedStorageType/UnmanagedStorageType.ql

@@ -0,0 +1,3 @@
+class A {
+    unowned(unsafe) var x: A?
+}