Kotlin: Handle enums better when generating dbscheme

java/kotlin-extractor/generate_dbscheme.py

@@ -13,9 +13,54 @@ with open('../ql/src/config/semmlecode.dbscheme', 'r') as f:
 dbscheme = re.sub(r'/\*.*?\*/', '', dbscheme, flags=re.DOTALL)
 dbscheme = re.sub(r'//[^\r\n]*/', '', dbscheme)
 
+enums = {}
 type_aliases = {}
 type_hierarchy = {}
 
+def genTable(kt, relname, body, enum = None, kind = None, num = None, typ = None):
+    kt.write('fun TrapWriter.write' + upperFirst(relname))
+    if kind is not None:
+        kt.write('_' + typ)
+    kt.write('(')
+    for colname, db_type in re.findall('(\S+)\s*:\s*([^\s,]+)', body):
+        if colname != kind:
+            kt.write(colname + ': ')
+            if db_type == 'int':
+                # TODO: Do something better if the column is a 'case'
+                kt.write('Int')
+            elif db_type == 'float':
+                kt.write('Double')
+            elif db_type == 'string':
+                kt.write('String')
+            elif db_type == 'date':
+                kt.write('String')
+            elif db_type == 'boolean':
+                kt.write('Boolean')
+            elif db_type[0] == '@':
+                label = db_type[1:]
+                if label == enum:
+                    label = typ
+                kt.write('Label<out Db' + upperFirst(label) + '>')
+            else:
+                raise Exception('Bad db_type: ' + db_type)
+            kt.write(', ')
+    kt.write(') {\n')
+    kt.write('    this.writeTrap("' + relname + '(')
+    comma = ''
+    for colname, db_type in re.findall('(\S+)\s*:\s*([^\s,]+)', body):
+        kt.write(comma)
+        if colname == kind:
+            kt.write(str(num))
+        elif db_type == 'string' or db_type == 'date':
+            kt.write('\\"$' + colname + '\\"') # TODO: Escaping
+        else:
+            # TODO: Any reformatting or escaping necessary?
+            # e.g. float formats?
+            kt.write('$' + colname)
+        comma = ', '
+    kt.write(')\\n")\n')
+    kt.write('}\n')
+
 with open('src/main/kotlin/KotlinExtractorDbScheme.kt', 'w') as kt:
     kt.write('/* Generated by ' + sys.argv[0] + ': Do not edit manually. */\n')
     kt.write('package com.github.codeql\n')
@@ -29,10 +74,13 @@ with open('src/main/kotlin/KotlinExtractorDbScheme.kt', 'w') as kt:
     for name, kind, body in re.findall(r'case\s+@([^.\s]*)\.([^.\s]*)\s+of\b(.*?);',
                                        dbscheme,
                                        flags=re.DOTALL):
+        mapping = []
         for num, typ in re.findall(r'(\d+)\s*=\s*@(\S+)', body):
             s = type_hierarchy.get(typ, set())
             s.add(name)
             type_hierarchy[typ] = s
+            mapping.append((int(num), typ))
+        enums[name] = (kind, mapping)
 
     # unions
     for name, unions in re.findall(r'@(\w+)\s*=\s*(@\w+(?:\s*\|\s*@\w+)*)',
@@ -52,41 +100,17 @@ with open('src/main/kotlin/KotlinExtractorDbScheme.kt', 'w') as kt:
     for relname, body in re.findall('\n([\w_]+)(\([^)]*\))',
                                     dbscheme,
                                     flags=re.DOTALL):
+        enum = None
         for db_type in re.findall(':\s*@([^\s,]+)\s*(?:,|$)', body):
             type_hierarchy[db_type] = type_hierarchy.get(db_type, set())
-        kt.write('fun TrapWriter.write' + upperFirst(relname) + '(')
-        for colname, db_type in re.findall('(\S+)\s*:\s*([^\s,]+)', body):
-            kt.write(colname + ': ')
-            if db_type == 'int':
-                # TODO: Do something better if the column is a 'case'
-                kt.write('Int')
-            elif db_type == 'float':
-                kt.write('Double')
-            elif db_type == 'string':
-                kt.write('String')
-            elif db_type == 'date':
-                kt.write('String')
-            elif db_type == 'boolean':
-                kt.write('Boolean')
-            elif db_type[0] == '@':
-                kt.write('Label<out Db' + upperFirst(db_type[1:]) + '>')
-            else:
-                raise Exception('Bad db_type: ' + db_type)
-            kt.write(', ')
-        kt.write(') {\n')
-        kt.write('    this.writeTrap("' + relname + '(')
-        comma = ''
-        for colname, db_type in re.findall('(\S+)\s*:\s*([^\s,]+)', body):
-            kt.write(comma)
-            if db_type == 'string' or db_type == 'date':
-                kt.write('\\"$' + colname + '\\"') # TODO: Escaping
-            else:
-                # TODO: Any reformatting or escaping necessary?
-                # e.g. float formats?
-                kt.write('$' + colname)
-            comma = ', '
-        kt.write(')\\n")\n')
-        kt.write('}\n')
+            if db_type in enums:
+                enum = db_type
+        if enum is None:
+            genTable(kt, relname, body)
+        else:
+            (kind, mapping) = enums[enum]
+            for num, typ in mapping:
+                genTable(kt, relname, body, enum, kind, num, typ)
 
     for typ in sorted(type_hierarchy):
         if typ in type_aliases:

java/kotlin-extractor/src/main/kotlin/KotlinExtractorExtension.kt

@@ -223,8 +223,7 @@ class KotlinFileExtractor(val tw: TrapWriter) {
     fun extractBlockBody(b: IrBlockBody, callable: Label<out DbCallable>, parent: Label<out DbStmtparent>, idx: Int) {
         val id = tw.getFreshIdLabel<DbBlock>()
         val locId = tw.getLocation(b.startOffset, b.endOffset)
-        val kind = 0 // TODO: stmt kind for block from generated module
-        tw.writeStmts(id, kind, parent, idx, callable)
+        tw.writeStmts_block(id, parent, idx, callable)
         tw.writeHasLocation(id, locId)
         for((sIdx, stmt) in b.statements.withIndex()) {
             extractStatement(stmt, callable, id, sIdx)
@@ -236,8 +235,7 @@ class KotlinFileExtractor(val tw: TrapWriter) {
             is IrReturn -> {
                 val id = tw.getFreshIdLabel<DbReturnstmt>()
                 val locId = tw.getLocation(s.startOffset, s.endOffset)
-                val kind = 9 // TODO: stmt kind for return from generated module
-                tw.writeStmts(id, kind, parent, idx, callable)
+                tw.writeStmts_returnstmt(id, parent, idx, callable)
                 tw.writeHasLocation(id, locId)
                 extractExpression(s.value, id, 0)
             }
@@ -252,11 +250,10 @@ class KotlinFileExtractor(val tw: TrapWriter) {
                     val left = e.dispatchReceiver
                     val right = e.getValueArgument(0)
                     if(left != null && right != null) {
-                        val kind = 27 // TODO: expr kind for addexpr from generated module
                         val id = tw.getFreshIdLabel<DbAddexpr>()
                         val typeId = useType(e.type)
                         val locId = tw.getLocation(e.startOffset, e.endOffset)
-                        tw.writeExprs(id, kind, typeId, parent, idx)
+                        tw.writeExprs_addexpr(id, typeId, parent, idx)
                         tw.writeHasLocation(id, locId)
                         extractExpression(left, id, 0)
                         extractExpression(right, id, 1)
@@ -269,11 +266,10 @@ class KotlinFileExtractor(val tw: TrapWriter) {
             }
             is IrConst<*> -> {
                 val v = e.value as Int
-                val kind = 17 // TODO: expr kind for integerliteral from generated module
                 val id = tw.getFreshIdLabel<DbIntegerliteral>()
                 val typeId = useType(e.type)
                 val locId = tw.getLocation(e.startOffset, e.endOffset)
-                tw.writeExprs(id, kind, typeId, parent, idx)
+                tw.writeExprs_integerliteral(id, typeId, parent, idx)
                 tw.writeHasLocation(id, locId)
                 tw.writeNamestrings(v.toString(), v.toString(), id)
             }