Initial support for oneofs

3 files changed, 161 insertions, 18 deletions

diff --git a/generator/descriptor_pb.nim b/generator/descriptor_pb.nim
index 5d07365..2a734e8 100644
--- a/generator/descriptor_pb.nim
+++ b/generator/descriptor_pb.nim
@@ -100,6 +100,14 @@ const
                 typeName: "EnumDescriptorProto",
                 packed: false
             ),
+            FieldDesc(
+                name: "oneof_decl",
+                number: 8,
+                ftype: FieldType.Message,
+                label: FieldLabel.Repeated,
+                typeName: "OneofDescriptorProto",
+                packed: false
+            ),
         ]
     )
 
@@ -198,6 +206,14 @@ const
                 typeName: "FieldOptions",
                 packed: false
             ),
+            FieldDesc(
+                name: "oneof_index",
+                number: 9,
+                ftype: FieldType.Int32,
+                label: FieldLabel.Optional,
+                typeName: "",
+                packed: false
+            ),
         ]
     )
 
@@ -248,6 +264,20 @@ const
         ]
     )
 
+    OneofDescriptorProtoDesc = MessageDesc(
+        name: "OneofDescriptorProto",
+        fields: @[
+            FieldDesc(
+                name: "name",
+                number: 1,
+                ftype: FieldType.String,
+                label: FieldLabel.Optional,
+                typeName: "",
+                packed: false
+            ),
+        ]
+    )
+
 generateEnumType(FieldDescriptorProto_LabelDesc)
 generateEnumProcs(FieldDescriptorProto_LabelDesc)
 
@@ -266,6 +296,9 @@ generateMessageProcs(FieldOptionsDesc)
 generateMessageType(FieldDescriptorProtoDesc)
 generateMessageProcs(FieldDescriptorProtoDesc)
 
+generateMessageType(OneofDescriptorProtoDesc)
+generateMessageProcs(OneofDescriptorProtoDesc)
+
 generateMessageType(DescriptorProtoDesc)
 generateMessageProcs(DescriptorProtoDesc)
 
diff --git a/generator/protoc_gen_nim.nim b/generator/protoc_gen_nim.nim
index 8b38815..f35f1cc 100644
--- a/generator/protoc_gen_nim.nim
+++ b/generator/protoc_gen_nim.nim
@@ -27,10 +27,12 @@ type
         ftype: FieldType
         typeName: string
         packed: bool
+        oneofIdx: int
 
     Message = ref object
         names: Names
         fields: seq[Field]
+        oneofs: seq[string]
 
     ProcessedFile = ref object
         name: string
@@ -89,6 +91,11 @@ proc newField(desc: FieldDescriptorProto): Field =
     result.ftype = convertFieldType(desc.type)
     result.typeName = ""
     result.packed = desc.options.packed
+    result.oneofIdx =
+        if hasOneof_index(desc):
+            desc.oneof_index
+        else:
+            -1
 
     if result.ftype == FieldType.Message or result.ftype == FieldType.Enum:
         result.typeName = $initNamesFromTypeName(desc.type_name)
@@ -100,12 +107,16 @@ proc newMessage(names: Names, desc: DescriptorProto): Message =
 
     result.names = names
     result.fields = @[]
+    result.oneofs = @[]
 
     log(&"newMessage {$result.names}")
 
     for field in desc.field:
         add(result.fields, newField(field))
 
+    for oneof in desc.oneof_decl:
+        add(result.oneofs, oneof.name)
+
 proc newEnum(names: Names, desc: EnumDescriptorProto): Enum =
     new(result)
 
@@ -238,6 +249,7 @@ proc generateDesc(field: Field): string =
     addLine(result, &"label: FieldLabel.{field.label},", 16)
     addLine(result, &"typeName: \"{field.typeName}\",", 16)
     addLine(result, &"packed: {field.packed},", 16)
+    addLine(result, &"oneofIdx: {field.oneofIdx},", 16)
     addLine(result, "),", 12)
 
 proc generateDesc(message: Message): string =
@@ -247,7 +259,11 @@ proc generateDesc(message: Message): string =
     addLine(result, "fields: @[", 8)
     for field in message.fields:
         result &= generateDesc(field)
-    addLine(result, "]", 8)
+    addLine(result, "],", 8)
+    addLine(result, "oneofs: @[", 8)
+    for oneof in message.oneofs:
+        addLine(result, &"\"{oneof}\",", 12)
+    addLine(result, "],", 8)
     addLine(result, ")", 4)
 
 proc generateDesc(e: Enum): string =
diff --git a/src/protobuf/gen.nim b/src/protobuf/gen.nim
index fd3b3eb..c21565e 100644
--- a/src/protobuf/gen.nim
+++ b/src/protobuf/gen.nim
@@ -7,6 +7,7 @@ type
     MessageDesc* = object
         name*: string
         fields*: seq[FieldDesc]
+        oneofs*: seq[string]
 
     FieldLabel* {.pure.} = enum
         Optional = 1
@@ -20,6 +21,7 @@ type
         label*: FieldLabel
         typeName*: string
         packed*: bool
+        oneofIdx*: int
 
     EnumDesc* = object
         name*: string
@@ -137,21 +139,81 @@ proc defaultValue(field: NimNode): NimNode =
 proc wiretype(field: NimNode): WireType =
     result = wiretype(getFieldType(field))
 
-proc fieldInitializer(objname: string, field: NimNode): NimNode =
+# TODO: maybe not the best name for this
+proc getFieldNameAST(objname: NimNode, field: NimNode, oneof: string): NimNode =
+    result =
+        if oneof != "":
+            newDotExpr(newDotExpr(objname, ident(oneof)), ident(getFieldName(field)))
+        else:
+            newDotExpr(objname, ident(getFieldName(field)))
+
+proc fieldInitializer(objname: NimNode, field: NimNode, oneof: string): NimNode =
     result = nnkAsgn.newTree(
-        nnkDotExpr.newTree(
-            newIdentNode(objname),
-            newIdentNode(getFieldName(field))
-        ),
+        getFieldNameAST(objname, field, oneof),
         defaultValue(field)
     )
 
+proc oneofIndex(field: NimNode): int =
+    let node = findColonExpr(field, "oneofIdx")
+    if node == nil:
+        result = -1
+    else:
+        result = int(node[1].intVal)
+
+proc oneofName(message, field: NimNode): string =
+    let index = oneofIndex(field)
+
+    if index == -1:
+        return ""
+
+    let oneofs = findColonExpr(message, "oneofs")[1]
+
+    result = $oneofs[index]
+
+iterator oneofFields(message: NimNode, index: int): NimNode =
+    if index != -1:
+        for field in fields(message):
+            if oneofIndex(field) == index:
+                yield field
+
+proc generateOneofFields*(desc: NimNode, typeSection: NimNode) =
+    let
+        oneofs = findColonExpr(desc, "oneofs")[1]
+        messageName = getMessageName(desc)
+
+    for index, oneof in oneofs:
+        let reclist = nnkRecList.newTree()
+
+        for field in oneofFields(desc, index):
+            let ftype = getFullFieldType(field)
+            let name = ident(getFieldName(field))
+
+            add(reclist, newIdentDefs(postfix(name, "*"), ftype))
+
+        let typedef = nnkTypeDef.newTree(
+            nnkPragmaExpr.newTree(
+                postfix(ident(messageName & $oneof), "*"),
+                nnkPragma.newTree(
+                    ident("union")
+                )
+            ),
+            newEmptyNode(),
+            nnkObjectTy.newTree(
+                newEmptyNode(),
+                newEmptyNode(),
+                reclist
+            )
+        )
+
+        add(typeSection, typedef)
+
 macro generateMessageType*(desc: typed): typed =
     let
         impl = getImpl(symbol(desc))
         typeSection = nnkTypeSection.newTree()
         typedef = nnkTypeDef.newTree()
         reclist = nnkRecList.newTree()
+        oneofs = findColonExpr(impl, "oneofs")[1]
 
     let name = getMessageName(impl)
 
@@ -168,11 +230,18 @@ macro generateMessageType*(desc: typed): typed =
     for field in fields(impl):
         let ftype = getFullFieldType(field)
         let name = ident(getFieldName(field))
-        add(reclist, newIdentDefs(postfix(name, "*"), ftype))
+        if oneofIndex(field) == -1:
+            add(reclist, newIdentDefs(postfix(name, "*"), ftype))
+
+    for oneof in oneofs:
+        add(reclist, newIdentDefs(postfix(ident($oneof), "*"),
+            ident(name & $oneof)))
 
     add(reclist, nnkIdentDefs.newTree(
         ident("hasField"), ident("IntSet"), newEmptyNode()))
 
+    generateOneofFields(impl, typeSection)
+
     result = newStmtList()
     add(result, typeSection)
 
@@ -180,14 +249,17 @@ macro generateMessageType*(desc: typed): typed =
         hint(repr(result))
 
 proc generateNewMessageProc(desc: NimNode): NimNode =
-    let body = newStmtList(
-        newCall(ident("new"), ident("result"))
-    )
+    let
+        body = newStmtList(
+            newCall(ident("new"), ident("result"))
+        )
+        resultId = ident("result")
 
     for field in fields(desc):
-        add(body, fieldInitializer("result", field))
+        let oneofName = oneofName(desc, field)
+        add(body, fieldInitializer(resultId, field, oneofName))
 
-    add(body, newAssignment(newDotExpr(ident("result"), ident("hasField")),
+    add(body, newAssignment(newDotExpr(resultId, ident("hasField")),
         newCall(ident("initIntSet"))))
 
     result = newProc(postfix(ident("new" & getMessageName(desc)), "*"),
@@ -203,7 +275,7 @@ proc fieldProcIdent(prefix: string, field: NimNode): NimNode =
 proc generateClearFieldProc(desc, field: NimNode): NimNode =
     let
         messageId = ident("message")
-        fname = newDotExpr(messageId, ident(getFieldName(field)))
+        fname = getFieldNameAST(messageId, field, oneofName(desc, field))
         defvalue = defaultValue(field)
         hasField = newDotExpr(messageId, ident("hasField"))
         number = getFieldNumber(field)
@@ -215,6 +287,17 @@ proc generateClearFieldProc(desc, field: NimNode): NimNode =
             `fname` = `defvalue`
             excl(`hasfield`, `number`)
 
+    # When clearing a field that is contained in a oneof, we should also clear
+    # the other fields.
+    for sibling in oneofFields(desc, oneofIndex(field)):
+        if sibling == field:
+            continue
+        let
+            number = getFieldNumber(sibling)
+            exclNode = quote do:
+                excl(`hasField`, `number`)
+        add(body(result), exclNode)
+
 proc generateHasFieldProc(desc, field: NimNode): NimNode =
     let
         messageId = ident("message")
@@ -233,7 +316,7 @@ proc generateSetFieldProc(desc, field: NimNode): NimNode =
         hasField = newDotExpr(messageId, ident("hasField"))
         number = getFieldNumber(field)
         valueId = ident("value")
-        fname = newDotExpr(messageId, ident(getFieldName(field)))
+        fname = getFieldNameAST(messageId, field, oneofName(desc, field))
         procName = fieldProcIdent("set", field)
         mtype = ident(getMessageName(desc))
         ftype = getFullFieldType(field)
@@ -243,6 +326,16 @@ proc generateSetFieldProc(desc, field: NimNode): NimNode =
             `fname` = `valueId`
             incl(`hasfield`, `number`)
 
+    # When setting a field that is in a oneof, we need to unset the other fields
+    for sibling in oneofFields(desc, oneofIndex(field)):
+        if sibling == field:
+            continue
+        let
+            number = getFieldNumber(sibling)
+            exclNode = quote do:
+                excl(`hasField`, `number`)
+        add(body(result), exclNode)
+
 proc generateAddToFieldProc(desc, field: NimNode): NimNode =
     let
         procName = fieldProcIdent("add", field)
@@ -262,13 +355,14 @@ proc generateAddToFieldProc(desc, field: NimNode): NimNode =
 proc ident(wt: WireType): NimNode =
     result = newDotExpr(ident("WireType"), ident($wt))
 
-proc genWriteField(field: NimNode): NimNode =
+proc genWriteField(message, field: NimNode): NimNode =
     result = newStmtList()
 
     let
         number = getFieldNumber(field)
         writer = ident("write" & getFieldTypeAsString(field))
-        fname = newDotExpr(ident("message"), ident(getFieldName(field)))
+        messageId = ident("message")
+        fname = getFieldNameAST(messageId, field, oneofName(message, field))
         wiretype = ident(wiretype(field))
         sizeproc = ident("sizeOf" & getFieldTypeAsString(field))
         hasproc = ident(fieldProcName("has", field))
@@ -310,7 +404,7 @@ proc generateWriteMessageProc(desc: NimNode): NimNode =
         sizeproc = postfix(ident("sizeOf" & getMessageName(desc)), "*")
 
     for field in fields(desc):
-        add(body, genWriteField(field))
+        add(body, genWriteField(desc, field))
 
     result = quote do:
         proc `sizeproc`(`messageId`: `mtype`): uint64
@@ -410,7 +504,7 @@ proc generateSizeOfMessageProc(desc: NimNode): NimNode =
         let
             hasproc = ident(fieldProcName("has", field))
             sizeofproc = ident("sizeOf" & getFieldTypeAsString(field))
-            fname = newDotExpr(messageId, ident(getFieldName(field)))
+            fname = getFieldNameAST(messageId, field, oneofName(desc, field))
             number = getFieldNumber(field)
             wiretype = ident(wiretype(field))