Provide searcher suggestions for compatible types (#1613)

engine/language-server/src/main/scala/org/enso/languageserver/search/SuggestionsHandler.scala

@@ -1,7 +1,7 @@
 package org.enso.languageserver.search
 
 import akka.actor.{Actor, ActorLogging, ActorRef, Props, Stash}
-import akka.pattern.pipe
+import akka.pattern.{ask, pipe}
 import org.enso.languageserver.capability.CapabilityProtocol.{
   AcquireCapability,
   CapabilityAcquired,
@@ -26,6 +26,7 @@ import org.enso.languageserver.session.SessionRouter.DeliverToJsonController
 import org.enso.languageserver.util.UnhandledLogging
 import org.enso.pkg.PackageManager
 import org.enso.polyglot.Suggestion
+import org.enso.polyglot.data.TypeGraph
 import org.enso.polyglot.runtime.Runtime.Api
 import org.enso.searcher.data.QueryResult
 import org.enso.searcher.{FileVersionsRepo, SuggestionsRepo}
@@ -98,6 +99,10 @@ final class SuggestionsHandler(
     context.system.eventStream
       .subscribe(self, InitializedEvent.SuggestionsRepoInitialized.getClass)
 
+    runtimeConnector
+      .ask(Api.Request(Api.GetTypeGraphRequest()))(timeout, self)
+      .pipeTo(self)
+
     config.contentRoots.foreach { case (_, contentRoot) =>
       PackageManager.Default
         .fromDirectory(contentRoot)
@@ -122,23 +127,29 @@ final class SuggestionsHandler(
           Some(InitializedEvent.SuggestionsRepoInitialized)
         )
       )
+    case Api.GetTypeGraphResponse(g) =>
+      tryInitialize(init.copy(typeGraph = Some(g)))
     case _ => stash()
   }
 
-  def initialized(projectName: String, clients: Set[ClientId]): Receive = {
+  def initialized(
+    projectName: String,
+    graph: TypeGraph,
+    clients: Set[ClientId]
+  ): Receive = {
     case AcquireCapability(
           client,
           CapabilityRegistration(ReceivesSuggestionsDatabaseUpdates())
         ) =>
       sender() ! CapabilityAcquired
-      context.become(initialized(projectName, clients + client.clientId))
+      context.become(initialized(projectName, graph, clients + client.clientId))
 
     case ReleaseCapability(
           client,
           CapabilityRegistration(ReceivesSuggestionsDatabaseUpdates())
         ) =>
       sender() ! CapabilityReleased
-      context.become(initialized(projectName, clients - client.clientId))
+      context.become(initialized(projectName, graph, clients - client.clientId))
 
     case msg: Api.SuggestionsDatabaseModuleUpdateNotification =>
       val isVersionChanged =
@@ -216,14 +227,16 @@ final class SuggestionsHandler(
         .pipeTo(sender())
 
     case Completion(path, pos, selfType, returnType, tags) =>
+      val selfTypes =
+        selfType.toList.flatMap(ty => (graph.getParents(ty) + ty).toSeq)
       getModuleName(projectName, path)
         .fold(
           Future.successful,
           module =>
             suggestionsRepo
               .search(
                 Some(module),
-                selfType,
+                selfTypes,
                 returnType,
                 tags.map(_.map(SuggestionKind.toSuggestion)),
                 Some(toPosition(pos))
@@ -295,18 +308,19 @@ final class SuggestionsHandler(
         .pipeTo(self)
 
     case ProjectNameUpdated(name) =>
-      context.become(initialized(name, clients))
+      context.become(initialized(name, graph, clients))
   }
 
   /** Transition the initialization process.
     *
     * @param state current initialization state
     */
   private def tryInitialize(state: SuggestionsHandler.Initialization): Unit = {
-    state.initialized.fold(context.become(initializing(state))) { name =>
-      log.debug("Initialized")
-      context.become(initialized(name, Set()))
-      unstashAll()
+    state.initialized.fold(context.become(initializing(state))) {
+      case (name, graph) =>
+        log.debug("Initialized")
+        context.become(initialized(name, graph, Set()))
+        unstashAll()
     }
   }
 
@@ -443,21 +457,25 @@ object SuggestionsHandler {
     *
     * @param project the project name
     * @param suggestions the initialization event of the suggestions repo
+    * @param typeGraph the Enso type hierarchy
     */
   private case class Initialization(
-    project: Option[String]                                               = None,
-    suggestions: Option[InitializedEvent.SuggestionsRepoInitialized.type] = None
+    project: Option[String] = None,
+    suggestions: Option[InitializedEvent.SuggestionsRepoInitialized.type] =
+      None,
+    typeGraph: Option[TypeGraph] = None
   ) {
 
     /** Check if all the components are initialized.
       *
       * @return the project name
       */
-    def initialized: Option[String] =
+    def initialized: Option[(String, TypeGraph)] =
       for {
-        _    <- suggestions
-        name <- project
-      } yield name
+        _     <- suggestions
+        name  <- project
+        graph <- typeGraph
+      } yield (name, graph)
   }
 
   /** Creates a configuration object used to create a [[SuggestionsHandler]].
@@ -484,5 +502,4 @@ object SuggestionsHandler {
         runtimeConnector
       )
     )
-
 }

engine/language-server/src/test/scala/org/enso/languageserver/search/Suggestions.scala

@@ -7,57 +7,95 @@ import org.enso.polyglot.Suggestion
 /** Suggestion instances used in tests. */
 object Suggestions {
 
-  val atom: Suggestion.Atom =
-    Suggestion.Atom(
-      externalId    = None,
-      module        = "Test.Main",
-      name          = "MyType",
-      arguments     = Vector(Suggestion.Argument("a", "Any", false, false, None)),
-      returnType    = "MyAtom",
-      documentation = None
-    )
-
-  val method: Suggestion.Method =
-    Suggestion.Method(
-      externalId =
-        Some(UUID.fromString("ea9d7734-26a7-4f65-9dd9-c648eaf57d63")),
-      module = "Test.Main",
-      name   = "foo",
-      arguments = Vector(
-        Suggestion.Argument("this", "MyType", false, false, None),
-        Suggestion.Argument("foo", "Number", false, true, Some("42"))
-      ),
-      selfType      = "MyType",
-      returnType    = "Number",
-      documentation = Some("Lovely")
-    )
-
-  val function: Suggestion.Function =
-    Suggestion.Function(
-      externalId =
-        Some(UUID.fromString("78d452ce-ed48-48f1-b4f2-b7f45f8dff89")),
-      module = "Test.Main",
-      name   = "print",
-      arguments = Vector(
-        Suggestion.Argument("a", "Any", false, false, None),
-        Suggestion.Argument("b", "Any", true, false, None),
-        Suggestion.Argument("c", "Any", false, true, Some("C"))
-      ),
-      returnType = "IO",
-      scope =
-        Suggestion.Scope(Suggestion.Position(1, 9), Suggestion.Position(1, 22))
-    )
-
-  val local: Suggestion.Local =
-    Suggestion.Local(
-      externalId =
-        Some(UUID.fromString("dc077227-d9b6-4620-9b51-792c2a69419d")),
-      module     = "Test.Main",
-      name       = "x",
-      returnType = "Number",
-      scope =
-        Suggestion.Scope(Suggestion.Position(21, 0), Suggestion.Position(89, 0))
-    )
-
-  val all = Seq(atom, method, function, local)
+  val atom: Suggestion.Atom = Suggestion.Atom(
+    externalId    = None,
+    module        = "Test.Main",
+    name          = "MyType",
+    arguments     = Vector(Suggestion.Argument("a", "Any", false, false, None)),
+    returnType    = "MyAtom",
+    documentation = None
+  )
+
+  val method: Suggestion.Method = Suggestion.Method(
+    externalId = Some(UUID.fromString("ea9d7734-26a7-4f65-9dd9-c648eaf57d63")),
+    module     = "Test.Main",
+    name       = "foo",
+    arguments = Vector(
+      Suggestion.Argument("this", "MyType", false, false, None),
+      Suggestion.Argument("foo", "Number", false, true, Some("42"))
+    ),
+    selfType      = "MyType",
+    returnType    = "Number",
+    documentation = Some("Lovely")
+  )
+
+  val function: Suggestion.Function = Suggestion.Function(
+    externalId = Some(UUID.fromString("78d452ce-ed48-48f1-b4f2-b7f45f8dff89")),
+    module     = "Test.Main",
+    name       = "print",
+    arguments = Vector(
+      Suggestion.Argument("a", "Any", false, false, None),
+      Suggestion.Argument("b", "Any", true, false, None),
+      Suggestion.Argument("c", "Any", false, true, Some("C"))
+    ),
+    returnType = "IO",
+    scope =
+      Suggestion.Scope(Suggestion.Position(1, 9), Suggestion.Position(1, 22))
+  )
+
+  val local: Suggestion.Local = Suggestion.Local(
+    externalId = Some(UUID.fromString("dc077227-d9b6-4620-9b51-792c2a69419d")),
+    module     = "Test.Main",
+    name       = "x",
+    returnType = "Number",
+    scope =
+      Suggestion.Scope(Suggestion.Position(21, 0), Suggestion.Position(89, 0))
+  )
+
+  val methodOnAny: Suggestion.Method = Suggestion.Method(
+    externalId = Some(UUID.fromString("6cfe1538-5df7-42e4-bf55-64f8ac2ededa")),
+    module     = "Standard.Base.Data.Any.Extensions",
+    name       = "<<",
+    arguments = Vector(
+      Suggestion.Argument("this", "Any", false, false, None),
+      Suggestion.Argument("that", "Any", false, false, None)
+    ),
+    selfType      = "Any",
+    returnType    = "Any",
+    documentation = Some("Lovely")
+  )
+
+  val methodOnNumber: Suggestion.Method = Suggestion.Method(
+    externalId = Some(UUID.fromString("33b426aa-2f74-42c0-9032-1159b5386eac")),
+    module     = "Standard.Base.Data.Number.Extensions",
+    name       = "asin",
+    arguments = Vector(
+      Suggestion.Argument("this", "Number", false, false, None)
+    ),
+    selfType      = "Number",
+    returnType    = "Number",
+    documentation = None
+  )
+
+  val methodOnInteger: Suggestion.Method = Suggestion.Method(
+    externalId = Some(UUID.fromString("2849c0f0-3c27-44df-abcb-5c163dd7ac91")),
+    module     = "Builtins.Main",
+    name       = "+",
+    arguments = Vector(
+      Suggestion.Argument("that", "Number", false, false, None)
+    ),
+    selfType      = "Integer",
+    returnType    = "Number",
+    documentation = Some("Blah, blah")
+  )
+
+  val all = Seq(
+    atom,
+    method,
+    function,
+    local,
+    methodOnAny,
+    methodOnNumber,
+    methodOnInteger
+  )
 }