Skip to content

Commit

Permalink
improvement: Add atomix update on TrieMap
Browse files Browse the repository at this point in the history
  • Loading branch information
jkciesluk committed Dec 29, 2023
1 parent ccd13ba commit 276a02c
Show file tree
Hide file tree
Showing 2 changed files with 69 additions and 21 deletions.
41 changes: 41 additions & 0 deletions mtags/src/main/scala/scala/meta/internal/mtags/AtomicTrieMap.scala
Original file line number Diff line number Diff line change
@@ -0,0 +1,41 @@
package scala.meta.internal.mtags

import java.util.concurrent.ConcurrentHashMap

import scala.collection.concurrent.TrieMap

/**
* This class is a wrapper around TrieMap that provides atomic updateWith
*/
final class AtomicTrieMap[K, V] {
private val trieMap = new TrieMap[K, V]()
private val concurrentMap = new ConcurrentHashMap[K, V]

def get(key: K): Option[V] = trieMap.get(key)

def contains(key: K): Boolean = trieMap.contains(key)

def updateWith(key: K)(remappingFunc: Option[V] => Option[V]): Unit = {
val computeFunction = new java.util.function.BiFunction[K, V, V] {
override def apply(k: K, v: V): V = {
trieMap.get(key) match {
case Some(value) =>
remappingFunc(Some(value)) match {
case Some(newValue) =>
trieMap.update(key, newValue)
case None =>
trieMap.remove(key)
}
case None =>
remappingFunc(None).foreach(trieMap.update(key, _))
}
null.asInstanceOf[V]
}
}
concurrentMap.compute(key, computeFunction)
}
}

object AtomicTrieMap {
def empty[K, V]: AtomicTrieMap[K, V] = new AtomicTrieMap[K, V]
}
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,6 @@ import java.nio.CharBuffer
import java.util.logging.Level
import java.util.logging.Logger

import scala.collection.concurrent.TrieMap
import scala.util.Properties
import scala.util.control.NonFatal

Expand Down Expand Up @@ -35,8 +34,8 @@ final case class SymbolLocation(
* while definitions contains only symbols generated by ScalaMtags.
*/
class SymbolIndexBucket(
toplevels: TrieMap[String, Set[AbsolutePath]],
definitions: TrieMap[String, Set[SymbolLocation]],
toplevels: AtomicTrieMap[String, Set[AbsolutePath]],
definitions: AtomicTrieMap[String, Set[SymbolLocation]],
sourceJars: OpenClassLoader,
toIndexSource: AbsolutePath => AbsolutePath = identity,
mtags: Mtags,
Expand Down Expand Up @@ -84,8 +83,10 @@ class SymbolIndexBucket(
): Unit = {
if (sourceJars.addEntry(jar.toNIO)) {
symbols.foreach { case (sym, path) =>
val acc = toplevels.getOrElse(sym, Set.empty)
toplevels(sym) = acc + path
toplevels.updateWith(sym) {
case Some(acc) => Some(acc + path)
case None => Some(Set(path))
}
}
}
}
Expand All @@ -96,8 +97,10 @@ class SymbolIndexBucket(
): List[String] = {
val symbols = indexSource(source, dialect, sourceDirectory)
symbols.foreach { symbol =>
val acc = toplevels.getOrElse(symbol, Set.empty)
toplevels(symbol) = acc + source
toplevels.updateWith(symbol) {
case Some(acc) => Some(acc + source)
case None => Some(Set(source))
}
}
symbols
}
Expand Down Expand Up @@ -132,8 +135,10 @@ class SymbolIndexBucket(
toplevel: String
): Unit = {
if (source.isAmmoniteScript || !isTrivialToplevelSymbol(path, toplevel)) {
val acc = toplevels.getOrElse(toplevel, Set.empty)
toplevels(toplevel) = acc + source
toplevels.updateWith(toplevel) {
case Some(acc) => Some(acc + source)
case None => Some(Set(source))
}
}
}

Expand Down Expand Up @@ -220,20 +225,20 @@ class SymbolIndexBucket(
.map(_.map(_.path))
.getOrElse(Set.empty)).filter(_.exists)

toplevels.get(symbol.value) match {
case None => ()
toplevels.updateWith(symbol.value) {
case None => None
case Some(acc) =>
val updated = acc.filter(exists(_))
if (updated.isEmpty) toplevels.remove(symbol.value)
else toplevels(symbol.value) = updated
if (updated.isEmpty) None
else Some(updated)
}

definitions.get(symbol.value) match {
case None => ()
definitions.updateWith(symbol.value) {
case None => None
case Some(acc) =>
val updated = acc.filter(loc => exists(loc.path))
if (updated.isEmpty) definitions.remove(symbol.value)
else definitions(symbol.value) = updated
if (updated.isEmpty) None
else Some(updated)
}
}

Expand Down Expand Up @@ -272,8 +277,10 @@ class SymbolIndexBucket(
docs.documents.foreach { document =>
document.occurrences.foreach { occ =>
if (occ.symbol.isGlobal && occ.role.isDefinition) {
val acc = definitions.getOrElse(occ.symbol, Set.empty)
definitions.put(occ.symbol, acc + SymbolLocation(file, occ.range))
definitions.updateWith(occ.symbol) {
case Some(acc) => Some(acc + SymbolLocation(file, occ.range))
case None => Some(Set(SymbolLocation(file, occ.range)))
}
} else {
// do nothing, we only care about global symbol definitions.
}
Expand Down Expand Up @@ -339,8 +346,8 @@ object SymbolIndexBucket {
toIndexSource: AbsolutePath => AbsolutePath
): SymbolIndexBucket =
new SymbolIndexBucket(
TrieMap.empty,
TrieMap.empty,
AtomicTrieMap.empty,
AtomicTrieMap.empty,
new OpenClassLoader,
toIndexSource,
mtags,
Expand Down

0 comments on commit 276a02c

Please sign in to comment.