fix(lsp): Plugin throws OOM on big projects (#1134)

language-server/language-server-api/.js/src/main/scala/aqua/lsp/ResultHelper.scala

@@ -130,9 +130,9 @@ object ResultHelper extends Logging {
     CompilationResult(
       errors.toJSArray,
       warnings.toJSArray,
-      locationsToJs(lsp.variables.flatMap(v => v.allLocations)),
+      locationsToJs(lsp.variables.allLocations),
       importTokens,
-      tokensToJs(lsp.variables.map(_.definition))
+      tokensToJs(lsp.variables.definitions)
     )
   }
 }

language-server/language-server-api/src/main/scala/aqua/lsp/LspContext.scala

@@ -4,14 +4,16 @@ import aqua.helpers.data.PName
 import aqua.parser.lexer.{LiteralToken, NamedTypeToken, Token}
 import aqua.raw.{RawContext, RawPart}
 import aqua.semantics.header.Picker
-import aqua.semantics.rules.locations.LocationsState
-import aqua.semantics.rules.locations.{TokenLocation, VariableInfo}
+import aqua.semantics.rules.locations.LocationsState.*
+import aqua.semantics.rules.locations.{LocationsState, TokenLocation, VariableInfo, Variables}
 import aqua.semantics.{SemanticError, SemanticWarning}
 import aqua.types.{AbilityType, ArrowType, Type}
 
 import cats.syntax.monoid.*
+import cats.syntax.semigroup.*
 import cats.{Monoid, Semigroup}
 import monocle.Lens
+import scala.collection.immutable.ListMap
 
 // Context with info that necessary for language server
 case class LspContext[S[_]](
@@ -20,13 +22,13 @@ case class LspContext[S[_]](
   rootArrows: Map[String, ArrowType] = Map.empty[String, ArrowType],
   constants: Map[String, Type] = Map.empty[String, Type],
   // TODO: Can this field be refactored into LocationsState?
-  variables: List[VariableInfo[S]] = Nil,
+  variables: Variables[S] = Variables[S](),
   importTokens: List[LiteralToken[S]] = Nil,
   errors: List[SemanticError[S]] = Nil,
   warnings: List[SemanticWarning[S]] = Nil,
   importPaths: Map[String, String] = Map.empty
 ) {
-  lazy val allLocations: List[TokenLocation[S]] = variables.flatMap(_.allLocations)
+  lazy val allLocations: List[TokenLocation[S]] = variables.allLocations
 }
 
 object LspContext {
@@ -43,7 +45,7 @@ object LspContext {
         rootArrows = x.rootArrows ++ y.rootArrows,
         constants = x.constants ++ y.constants,
         importTokens = x.importTokens ++ y.importTokens,
-        variables = x.variables ++ y.variables,
+        variables = x.variables |+| y.variables,
         errors = x.errors ++ y.errors,
         warnings = x.warnings ++ y.warnings,
         importPaths = x.importPaths ++ y.importPaths
@@ -79,13 +81,15 @@ object LspContext {
 
     override def allNames(ctx: LspContext[S]): Set[String] = ctx.raw.allNames
 
-    override def setAbility(ctx: LspContext[S], name: String, ctxAb: LspContext[S]): LspContext[S] =
+    override def setAbility(
+      ctx: LspContext[S],
+      name: String,
+      ctxAb: LspContext[S]
+    ): LspContext[S] =
       ctx.copy(
         raw = ctx.raw.setAbility(name, ctxAb.raw),
-        variables = ctx.variables ++ ctxAb.variables.map(v =>
-          v.copy(definition =
-            v.definition.copy(name = AbilityType.fullName(name, v.definition.name))
-          )
+        variables = ctx.variables |+| ctxAb.variables.renameDefinitions(defName =>
+          AbilityType.fullName(name, defName)
         )
       )
 
@@ -118,13 +122,9 @@ object LspContext {
     ): Option[LspContext[S]] =
       // rename tokens from one context with prefix addition
       val newVariables = rename.map { renameStr =>
-        ctx.variables.map {
-          case v if v.definition.name.startsWith(name) =>
-            v.copy(definition =
-              v.definition.copy(name = v.definition.name.replaceFirst(v.definition.name, renameStr))
-            )
-
-          case kv => kv
+        ctx.variables.renameDefinitions {
+          case defName if defName.startsWith(name) =>
+            defName.replaceFirst(name, renameStr)
         }
       }.getOrElse(ctx.variables)
 

language-server/language-server-api/src/main/scala/aqua/lsp/LspSemantics.scala

@@ -8,9 +8,9 @@ import aqua.semantics.*
 import aqua.semantics.header.Picker.*
 import aqua.semantics.rules.locations.LocationsState
 
-import cats.data.{EitherT, NonEmptyChain, Writer}
-import cats.syntax.functor.*
+import cats.data.EitherT
 import cats.syntax.applicative.*
+import cats.syntax.semigroup.*
 import monocle.Lens
 import monocle.macros.GenLens
 
@@ -41,7 +41,7 @@ class LspSemantics[S[_]](
 
     val initState = rawState.copy(
       locations = rawState.locations.copy(
-        variables = rawState.locations.variables ++ withConstants.variables
+        variables = rawState.locations.variables |+| withConstants.variables
       )
     )
 

language-server/language-server-api/src/test/scala/aqua/lsp/AquaLSPSpec.scala

@@ -43,7 +43,7 @@ class AquaLSPSpec extends AnyFlatSpec with Matchers with Inside {
       } yield {
         val (defStart, defEnd) = defPos
         val (useStart, useEnd) = usePos
-        c.variables.exists { case VariableInfo(defI, occs) =>
+        c.variables.variables.values.flatten.exists { case VariableInfo(defI, occs) =>
           val defSpan = defI.token.unit._1
           if (defSpan.startIndex == defStart && defSpan.endIndex == defEnd) {
             occs.exists { useT =>
@@ -76,7 +76,7 @@ class AquaLSPSpec extends AnyFlatSpec with Matchers with Inside {
     ): Boolean = {
 
       getByPosition(code, checkName, position).exists { case (start, end) =>
-        val res = c.variables.exists { case VariableInfo(definition, _) =>
+        val res = c.variables.variables.iterator.flatMap(_._2).exists { case VariableInfo(definition, _) =>
           val span = definition.token.unit._1
           definition.name == fullName.getOrElse(
             checkName
@@ -85,8 +85,7 @@ class AquaLSPSpec extends AnyFlatSpec with Matchers with Inside {
 
         if (printFiltered)
           println(
-            c.variables
-              .map(_.definition)
+            c.variables.definitions
               .filter(v => v.name == fullName.getOrElse(checkName))
               .map { case DefinitionInfo(name, token, t) =>
                 val span = token.unit._1

semantics/src/main/scala/aqua/semantics/rules/locations/LocationsState.scala

@@ -4,37 +4,27 @@ import aqua.helpers.syntax.list.*
 import aqua.parser.lexer.Token
 
 import cats.kernel.Monoid
+import cats.syntax.semigroup.*
 import scribe.Logging
 
 case class LocationsState[S[_]](
-  variables: List[VariableInfo[S]] = Nil
+  variables: Variables[S] = Variables[S]()
 ) extends Logging {
 
-  lazy val allLocations: List[TokenLocation[S]] = variables.flatMap(_.allLocations)
+  lazy val allLocations: List[TokenLocation[S]] = variables.allLocations
 
-  def addDefinitions(newDefinitions: List[DefinitionInfo[S]]): LocationsState[S] =
-    copy(variables = newDefinitions.map(d => VariableInfo(d)) ++ variables)
+  def addDefinitions(newDefinitions: List[DefinitionInfo[S]]): LocationsState[S] = {
+    copy(variables = variables.addDefinitions(newDefinitions))
+  }
 
   def addDefinition(newDef: DefinitionInfo[S]): LocationsState[S] =
-    copy(variables = VariableInfo(newDef) +: variables)
-
-  private def addOccurrenceToFirst(
-    vars: List[VariableInfo[S]],
-    name: String,
-    token: Token[S]
-  ): List[VariableInfo[S]] = {
-    // TODO: this code lasts too long, but we can find errors in it.
-    // if (!vars.exists(_.definition.name == name))
-    //   logger.error(s"Unexpected. Cannot add occurrence for $name")
-
-    vars.updateFirst(_.definition.name == name, v => v.copy(occurrences = token +: v.occurrences))
-  }
+    copy(variables = variables.addDefinitions(newDef :: Nil))
 
   def addLocation(
     name: String,
     token: Token[S]
   ): LocationsState[S] =
-    copy(variables = addOccurrenceToFirst(variables, name, token))
+    copy(variables = variables.addOccurence(name, token))
 
   def addLocations(
     locations: List[(String, Token[S])]
@@ -47,11 +37,11 @@ case class LocationsState[S[_]](
 object LocationsState {
 
   given [S[_]]: Monoid[LocationsState[S]] with {
-    override def empty: LocationsState[S] = LocationsState()
+    override def empty: LocationsState[S] = LocationsState[S]()
 
     override def combine(x: LocationsState[S], y: LocationsState[S]): LocationsState[S] =
       LocationsState(
-        variables = x.variables ++ y.variables
+        variables = x.variables.combine(y.variables)
       )
   }
 }

semantics/src/main/scala/aqua/semantics/rules/locations/Variables.scala

@@ -0,0 +1,64 @@
+package aqua.semantics.rules.locations
+
+import aqua.helpers.syntax.list.*
+import aqua.parser.lexer.Token
+
+import cats.kernel.{Monoid, Semigroup}
+import cats.syntax.align.*
+
+case class Variables[S[_]](
+  variables: Map[String, List[VariableInfo[S]]] = Map.empty[String, List[VariableInfo[S]]]
+) {
+
+  def renameDefinitions(f: PartialFunction[String, String]): Variables[S] =
+    copy(variables = variables.map { case (k, v) =>
+      f.andThen { newName =>
+        newName -> v.map(vi => vi.copy(definition = vi.definition.copy(name = newName)))
+      }.orElse { _ =>
+        k -> v
+      }(k)
+    })
+
+  lazy val allLocations: List[TokenLocation[S]] =
+    variables.values.flatMap(_.flatMap(_.allLocations)).toList
+
+  lazy val definitions: List[DefinitionInfo[S]] =
+    variables.values.flatMap(_.map(_.definition)).toList
+
+  def addDefinitions(newDefinitions: List[DefinitionInfo[S]]): Variables[S] = {
+    copy(variables =
+      newDefinitions
+        .map(d => d.name -> List(VariableInfo(d)))
+        .toMap
+        .alignCombine(variables)
+    )
+  }
+
+  /**
+   * Add occurrance by name to the first (last added) definition.
+   */
+  def addOccurence(
+    name: String,
+    token: Token[S]
+  ): Variables[S] = {
+    copy(variables =
+      variables.updatedWith(name)(
+        _.map(
+          _.updateFirst(
+            _.definition.name == name,
+            v => v.copy(occurrences = token +: v.occurrences)
+          )
+        )
+      )
+    )
+  }
+}
+
+object Variables {
+
+  given [S[_]]: Semigroup[Variables[S]] with {
+
+    override def combine(x: Variables[S], y: Variables[S]): Variables[S] =
+      Variables(x.variables.alignCombine(y.variables).view.mapValues(_.distinct).toMap)
+  }
+}

utils/helpers/src/main/scala/aqua/helpers/syntax/list.scala

@@ -3,7 +3,9 @@ package aqua.helpers.syntax
 import scala.annotation.tailrec
 
 object list {
+
   extension [A](l: List[A]) {
+
     def updateFirst[B >: A](p: A => Boolean, f: A => B): List[B] = {
       @tailrec
       def update(left: List[B], right: List[A]): List[B] =