Sum type info in suggestions DB

engine/polyglot-api/src/main/scala/org/enso/polyglot/Suggestion.scala

@@ -40,7 +40,7 @@ sealed trait Suggestion extends ToLogString {
   def externalId: Option[Suggestion.ExternalId]
   def module:     String
   def name:       String
-  def returnType: String
+  def returnType: Suggestion.TypeRep
 }
 
 object Suggestion {
@@ -110,6 +110,44 @@ object Suggestion {
       }
   }
 
+  @JsonTypeInfo(use = JsonTypeInfo.Id.NAME, property = "type")
+  @JsonSubTypes(
+    Array(
+      new JsonSubTypes.Type(
+        value = classOf[TypeRep.Legacy],
+        name  = "typeRepLegacy"
+      ),
+      new JsonSubTypes.Type(
+        value = classOf[TypeRep.Sum],
+        name  = "typeRepSum"
+      ),
+      new JsonSubTypes.Type(
+        value = classOf[TypeRep.Cons],
+        name  = "typeRepCons"
+      )
+    )
+  )
+  sealed trait TypeRep {
+    def toLegacy: TypeRep.Legacy
+  }
+
+  object TypeRep {
+    implicit def asLegacyRep(contents: String): TypeRep = Legacy(contents)
+
+    case class Legacy(contents: String) extends TypeRep {
+      override def toLegacy: Legacy = this
+      override def toString: String = contents
+    }
+
+    case class Sum(name: String, variants: Seq[TypeRep]) extends TypeRep {
+      override def toLegacy: Legacy = Legacy(this.name)
+    }
+
+    case class Cons(name: String) extends TypeRep {
+      override def toLegacy: Legacy = Legacy(this.name)
+    }
+  }
+
   /** An argument of an atom or a function.
     *
     * @param name the argument name
@@ -120,7 +158,7 @@ object Suggestion {
     */
   case class Argument(
     name: String,
-    reprType: String,
+    reprType: TypeRep,
     isSuspended: Boolean,
     hasDefault: Boolean,
     defaultValue: Option[String]
@@ -173,7 +211,7 @@ object Suggestion {
     override def externalId: Option[ExternalId] =
       None
 
-    override def returnType: String =
+    override def returnType: Suggestion.TypeRep =
       module
 
     /** @inheritdoc */
@@ -199,7 +237,7 @@ object Suggestion {
     module: String,
     name: String,
     arguments: Seq[Argument],
-    returnType: String,
+    returnType: Suggestion.TypeRep,
     documentation: Option[String],
     documentationHtml: Option[String]               = None,
     documentationSections: Option[List[DocSection]] = None,
@@ -238,7 +276,7 @@ object Suggestion {
     name: String,
     arguments: Seq[Argument],
     selfType: String,
-    returnType: String,
+    returnType: TypeRep,
     documentation: Option[String],
     documentationHtml: Option[String]               = None,
     documentationSections: Option[List[DocSection]] = None,
@@ -275,7 +313,7 @@ object Suggestion {
     module: String,
     arguments: Seq[Argument],
     sourceType: String,
-    returnType: String,
+    returnType: Suggestion.TypeRep,
     documentation: Option[String],
     documentationHtml: Option[String]               = None,
     documentationSections: Option[List[DocSection]] = None,
@@ -312,7 +350,7 @@ object Suggestion {
     module: String,
     name: String,
     arguments: Seq[Argument],
-    returnType: String,
+    returnType: Suggestion.TypeRep,
     scope: Scope
   ) extends Suggestion
       with ToLogString {
@@ -341,7 +379,7 @@ object Suggestion {
     externalId: Option[ExternalId],
     module: String,
     name: String,
-    returnType: String,
+    returnType: Suggestion.TypeRep,
     scope: Scope
   ) extends Suggestion {
 

engine/polyglot-api/src/main/scala/org/enso/polyglot/data/Tree.scala

@@ -30,6 +30,10 @@ sealed trait Tree[+A] {
   final def map[B](f: A => B): Tree[B] =
     Tree.map(this)(f)
 
+  @JsonIgnore
+  final def collectFirst[B](f: PartialFunction[A, B]): Option[B] =
+    Tree.collectFirst(this)(f)
+
   /** Selects all elements which satisfy a predicate.
     *
     * @param p the predicate used to test elements
@@ -208,6 +212,22 @@ object Tree {
         Node(f(a), cx.map(mapNode(_)(f)))
     }
 
+  private def collectFirst[A, B](
+    tree: Tree[A]
+  )(f: PartialFunction[A, B]): Option[B] = {
+    def collectFirstNode(node: Tree.Node[A]): Option[B] = {
+      if (f.isDefinedAt(node.element)) {
+        Some(f(node.element))
+      } else {
+        node.children.collectFirst((collectFirstNode _).unlift)
+      }
+    }
+    tree match {
+      case Root(c)    => c.collectFirst((collectFirstNode _).unlift)
+      case n: Node[A] => collectFirstNode(n)
+    }
+  }
+
   private def filter[A](tree: Tree[A])(p: A => Boolean): Tree[A] = {
     def filterNode(node: Tree.Node[A]): Option[Tree.Node[A]] = {
       val childrenFiltered = node.children.flatMap(filterNode)

engine/runtime/src/main/scala/org/enso/compiler/codegen/RuntimeStubsGenerator.scala

@@ -21,7 +21,7 @@ class RuntimeStubsGenerator() {
       BindingAnalysis,
       "Non-parsed module used in stubs generator"
     )
-    localBindings.types.foreach { tp =>
+    localBindings.constructors.foreach { tp =>
       val constructor = new AtomConstructor(tp.name, scope)
       scope.registerConstructor(constructor)
     }

engine/runtime/src/main/scala/org/enso/compiler/context/SuggestionBuilder.scala

@@ -336,13 +336,22 @@ final class SuggestionBuilder[A: IndexedSource](val source: A) {
     resolvedName match {
       case BindingsMap.ResolvedModule(module) =>
         Some(module.getName)
-      case BindingsMap.ResolvedConstructor(module, cons) =>
-        Some(module.getName.createChild(cons.name))
+      case cons: BindingsMap.ResolvedConstructor =>
+        Some(cons.qualifiedName)
       case _ =>
         None
     }
   }
 
+  private def buildResolvedUnionTypeName(
+    resolvedName: BindingsMap.ResolvedTypeName
+  ): Option[TypeArg] = resolvedName match {
+    case tp: BindingsMap.ResolvedType =>
+      Some(TypeArg.Sum(tp.qualifiedName, tp.getVariants.map(_.qualifiedName)))
+    case other: BindingsMap.ResolvedName =>
+      buildResolvedTypeName(other).map(TypeArg.Value)
+  }
+
   /** Build type signature from the ir metadata.
     *
     * @param typeSignature the type signature metadata
@@ -380,19 +389,18 @@ final class SuggestionBuilder[A: IndexedSource](val source: A) {
           args :++ arg
         case IR.Function.Lambda(List(targ), body, _, _, _, _) =>
           val typeName = targ.name.name
-          val qualifiedTypeName = resolveTypeName(bindings, typeName)
-            .getOrElse(QualifiedName.simpleName(typeName))
-          val tdef = TypeArg.Value(qualifiedTypeName)
+          val tdef = resolveTypeName(bindings, typeName)
+            .getOrElse(TypeArg.Value(QualifiedName.simpleName(typeName)))
           go(body, args :+ tdef)
         case IR.Application.Prefix(tfun, targs, _, _, _, _) =>
           val appFunction = go(tfun, Vector()).head
           val appArgs     = targs.flatMap(arg => go(arg.value, Vector()))
           args :+ TypeArg.Application(appFunction, appArgs.toVector)
         case tname: IR.Name =>
           val typeName = tname.name
-          val qualifiedTypeName = resolveTypeName(bindings, typeName)
-            .getOrElse(QualifiedName.simpleName(typeName))
-          args :+ TypeArg.Value(qualifiedTypeName)
+          val tdef = resolveTypeName(bindings, typeName)
+            .getOrElse(TypeArg.Value(QualifiedName.simpleName(typeName)))
+          args :+ tdef
         case _ =>
           args
       }
@@ -415,10 +423,10 @@ final class SuggestionBuilder[A: IndexedSource](val source: A) {
   private def resolveTypeName(
     bindings: Option[BindingAnalysis.Metadata],
     name: String
-  ): Option[QualifiedName] = {
+  ): Option[TypeArg] = {
     bindings
-      .flatMap(_.resolveUppercaseName(name).toOption)
-      .flatMap(buildResolvedTypeName)
+      .flatMap(_.resolveTypeName(name).toOption)
+      .flatMap(buildResolvedUnionTypeName)
   }
 
   /** Build arguments of a method.
@@ -529,8 +537,8 @@ final class SuggestionBuilder[A: IndexedSource](val source: A) {
     * @param targ the type argument
     * @return the name of type argument
     */
-  private def buildTypeArgumentName(targ: TypeArg): String = {
-    def go(targ: TypeArg, level: Int): String =
+  private def buildTypeArgumentName(targ: TypeArg): Suggestion.TypeRep = {
+    def go(targ: TypeArg, level: Int): Suggestion.TypeRep =
       targ match {
         case TypeArg.Value(name) => name.toString
         case TypeArg.Function(Vector(typeArg)) =>
@@ -549,6 +557,11 @@ final class SuggestionBuilder[A: IndexedSource](val source: A) {
           val argsList = args.map(go(_, level + 1)).mkString(" ")
           val typeName = s"$funText $argsList"
           if (level > 0) s"($typeName)" else typeName
+        case TypeArg.Sum(n, vs) =>
+          Suggestion.TypeRep.Sum(
+            n.toString,
+            vs.map(qn => Suggestion.TypeRep.Cons(qn.toString))
+          )
       }
 
     go(targ, 0)
@@ -573,7 +586,7 @@ final class SuggestionBuilder[A: IndexedSource](val source: A) {
     * @param typeDef the type definition
     * @return the type name
     */
-  private def buildReturnType(typeDef: Option[TypeArg]): String =
+  private def buildReturnType(typeDef: Option[TypeArg]): Suggestion.TypeRep =
     typeDef.map(buildTypeArgumentName).getOrElse(Any)
 
   /** Build argument default value from the expression.
@@ -639,6 +652,9 @@ object SuggestionBuilder {
   sealed private trait TypeArg
   private object TypeArg {
 
+    case class Sum(name: QualifiedName, variants: Seq[QualifiedName])
+        extends TypeArg
+
     /** Type with the name, like `A`.
       *
       * @param name the name of the type
@@ -672,6 +688,6 @@ object SuggestionBuilder {
 
   }
 
-  val Any: String = Constants.ANY
+  val Any: Suggestion.TypeRep = Constants.ANY
 
 }

engine/runtime/src/main/scala/org/enso/compiler/context/SuggestionDiff.scala

@@ -118,7 +118,7 @@ object SuggestionDiff {
       op = op.copy(name = Some(b.name))
     }
     if (a.reprType != b.reprType) {
-      op = op.copy(reprType = Some(b.reprType))
+      op = op.copy(reprType = Some(b.reprType.toLegacy.contents))
     }
     if (a.isSuspended != b.isSuspended) {
       op.copy(isSuspended = Some(b.isSuspended))
@@ -155,7 +155,7 @@ object SuggestionDiff {
       op = op.copy(arguments = Some(diffArguments(e1.arguments, e2.arguments)))
     }
     if (e1.returnType != e2.returnType) {
-      op = op.copy(returnType = Some(e2.returnType))
+      op = op.copy(returnType = Some(e2.returnType.toLegacy.contents))
     }
     if (e1.documentation != e2.documentation) {
       op = op.copy(documentation = Some(e2.documentation))
@@ -175,7 +175,7 @@ object SuggestionDiff {
       op = op.copy(arguments = Some(diffArguments(e1.arguments, e2.arguments)))
     }
     if (e1.returnType != e2.returnType) {
-      op = op.copy(returnType = Some(e2.returnType))
+      op = op.copy(returnType = Some(e2.returnType.toLegacy.contents))
     }
     if (e1.documentation != e2.documentation) {
       op = op.copy(documentation = Some(e2.documentation))
@@ -195,7 +195,7 @@ object SuggestionDiff {
       op = op.copy(arguments = Some(diffArguments(e1.arguments, e2.arguments)))
     }
     if (e1.returnType != e2.returnType) {
-      op = op.copy(returnType = Some(e2.returnType))
+      op = op.copy(returnType = Some(e2.returnType.toLegacy.contents))
     }
     if (e1.scope != e2.scope) {
       op = op.copy(scope = Some(e2.scope))
@@ -212,7 +212,7 @@ object SuggestionDiff {
       op = op.copy(externalId = Some(e2.externalId))
     }
     if (e1.returnType != e2.returnType) {
-      op = op.copy(returnType = Some(e2.returnType))
+      op = op.copy(returnType = Some(e2.returnType.toLegacy.contents))
     }
     if (e1.scope != e2.scope) {
       op = op.copy(scope = Some(e2.scope))