LruRef / LRU for cats-effect

_LruRef.scala

package teikametrics
import cats.Monad
import cats.effect.Sync
import cats.effect.concurrent.Ref
import cats.implicits._

object LruRef {

  /**
    * Build an immutable LRU key/value store that cannot grow larger than `maxSize`
    */
  def empty[F[_]: Sync, K, V](maxSize: Int): F[LruRef[F, K, V]] =
    Ref[F]
      .of(ImmutableLRU[K, V](maxSize))
      .map(new LruRef(_))

}

class LruRef[F[_]: Monad, K, V](lruRef: Ref[F, ImmutableLRU[K, V]]) {

  /** Adds the new kv to the LRU
    * @return The kv which was evicted to make space for the new input, if needed
    */
  def put(kv: (K, V)): F[Option[(K, V)]] =
    lruRef.modify { lru =>
      lru.put(kv).swap
    }

  /** @return The value at `k`, if present */
  def get(k: K): F[Option[V]] = lruRef.modify { lru =>
    lru.get(k).swap
  }

  /** @return The value at `k` prior to removal, if present */
  def remove(k: K): F[Option[V]] = lruRef.modify { lru =>
    lru.remove(k).swap
  }

  /** @return The number of keys present */
  def size: F[Int] = lruRef.get.map(_.size)

  /** @return The set of keys present */
  def keySet: F[Set[K]] = lruRef.get.map(_.keySet)
}

ImmutableLRU.scala

package teikametrics

import scala.collection.SortedMap

/**
  * Immutable implementation of an LRU cache.
  *
  * @author Twitter
  *
  *  Copy pasted from the version previously in twitter-util v19.1.0 at
  *  [[https://github.com/twitter/util/blob/1c748332580eb9b20f20a521429e1a6d79b1db82/util-collection/src/main/scala/com/twitter/util/ImmutableLRU.scala]]
  */
object ImmutableLRU {

  /**
    * Build an immutable LRU key/value store that cannot grow larger than `maxSize`
    */
  def apply[K, V](maxSize: Int): ImmutableLRU[K, V] =
    new ImmutableLRU(maxSize,
                     0,
                     Map.empty[K, (Long, V)],
                     SortedMap.empty[Long, K])
}

/**
  * An immutable key/value store that evicts the least recently accessed elements
  * to stay constrained in a maximum size bound.
  */
// "map" is the backing store used to hold key->(index,value)
// pairs. The index tracks the access time for a particular key. "ord"
// is used to determine the Least-Recently-Used key in "map" by taking
// the minimum index.
class ImmutableLRU[K, V] private (maxSize: Int,
                                  idx: Long,
                                  map: Map[K, (Long, V)],
                                  ord: SortedMap[Long, K]) {

  // Scala's SortedMap requires a key ordering; ImmutableLRU doesn't
  // care about pulling a minimum value out of the SortedMap, so the
  // following kOrd treats every value as equal.
  protected implicit val kOrd = new Ordering[K] { def compare(l: K, r: K) = 0 }

  /**
    * the number of entries in the cache
    */
  def size: Int = map.size

  /**
    * the `Set` of all keys in the LRU
    * @note accessing this set does not update the element LRU ordering
    */
  def keySet: Set[K] = map.keySet

  /** Alias for `put`, but discarding the old value */
  def +(kv: (K, V)): (Option[K], ImmutableLRU[K, V]) = this.put(kv) match {
    case (Some((k, _)), newLru) => Some(k) -> newLru
    case other @ (None, _) =>
      other.asInstanceOf[(Option[K], ImmutableLRU[K, V])]
  }

  /**
    * Build a new LRU containing the given key/value
    * @return a tuple with of the following two items:
    *          _1 represents the evicted entry (if the given lru is at the maximum size)
    *             or None if the lru is not at capacity yet.
    *          _2 is the new lru with the given key/value pair inserted.
    */
  def put(kv: (K, V)): (Option[(K, V)], ImmutableLRU[K, V]) = {
    val (key, value) = kv
    val newIdx = idx + 1
    val newMap = map + (key -> ((newIdx, value)))
    // Now update the ordered cache:
    val baseOrd = map.get(key).map { case (id, _) => ord - id }.getOrElse(ord)
    val ordWithNewKey = baseOrd + (newIdx -> key)
    // Do we need to remove an old key:
    val (evicts, finalMap, finalOrd) = if (ordWithNewKey.size > maxSize) {
      val (minIdx, eKey) = ordWithNewKey.min
      (Some(eKey -> map(eKey)._2), newMap - eKey, ordWithNewKey - minIdx)
    } else {
      (None, newMap, ordWithNewKey)
    }
    (evicts, new ImmutableLRU[K, V](maxSize, newIdx, finalMap, finalOrd))
  }

  /**
    * If the key is present in the cache, returns the pair of
    * Some(value) and the cache with the key's entry as the most recently accessed.
    * Else, returns None and the unmodified cache.
    */
  def get(k: K): (Option[V], ImmutableLRU[K, V]) = {
    val (optionalValue, lru) = remove(k)
    val newLru = optionalValue.map { v =>
      (lru + (k -> v))._2
    } getOrElse (lru)
    (optionalValue, newLru)
  }

  /**
    * If the key is present in the cache, returns the pair of
    * Some(value) and the cache with the key removed.
    * Else, returns None and the unmodified cache.
    */
  def remove(k: K): (Option[V], ImmutableLRU[K, V]) =
    map
      .get(k)
      .map {
        case (kidx, v) =>
          val newMap = map - k
          val newOrd = ord - kidx
          // Note we don't increase the idx on a remove, only on put:
          (Some(v), new ImmutableLRU[K, V](maxSize, idx, newMap, newOrd))
      }
      .getOrElse((None, this))

  override def toString = "ImmutableLRU(" + map.toList.mkString(",") + ")"
}

ImmutableLRUSpec.scala

package teikametrics
import org.scalacheck.Arbitrary.arbitrary
import org.scalacheck.{Arbitrary, Gen}
import org.scalatest.FunSuite
import org.scalatest.prop.GeneratorDrivenPropertyChecks

import scala.annotation.tailrec

/** Copied from twitter-util v19.1.0 at [[https://raw.githubusercontent.com/twitter/util/1c748332580eb9b20f20a521429e1a6d79b1db82/util-collection/src/test/scala/com/twitter/util/ImmutableLRUTest.scala]]
  * With adaptations for our test dependencies
  */
class ImmutableLRUSpec extends FunSuite with GeneratorDrivenPropertyChecks {

  // don't waste too much time testing this and keep things small
  implicit override val generatorDrivenConfig: PropertyCheckConfiguration =
    PropertyCheckConfiguration(minSuccessful = 5, minSize = 2, sizeRange = 8)

  test("ImmutableLRU insertion") {
    forAll(Gen.zip(Gen.identifier, arbitrary[Int])) {
      case (s: String, i: Int) =>
        val lru = ImmutableLRU[String, Int](50)

        // first-time insertion should not evict
        val (key, lru2) = lru + (s -> i)
        assert(key == None)

        // test get method
        val (key2, _) = lru2.get(s)
        assert(key2 == Some(i))
    }
  }

  // given a list of entries, build an LRU containing them
  @tailrec
  private def buildLRU[V](
      lru: ImmutableLRU[String, V],
      entries: List[(String, V)]
  ): ImmutableLRU[String, V] =
    entries match {
      case Nil          => lru
      case head :: tail => buildLRU((lru + head)._2, tail)
    }

  test("ImmutableLRU eviction") {
    forAll(LRUEntriesGenerator[Int]) { entries =>
      val lru = buildLRU(ImmutableLRU[String, Int](4), entries)
      assert(
        lru.keySet == entries
          .map(_._1)
          .slice(entries.size - 4, entries.size)
          .toSet)
    }
  }

  test("ImmutableLRU removal") {
    val gen = for {
      entries <- LRUEntriesGenerator[Double]
      entry <- Gen.oneOf(entries)
    } yield (entries, entry._1, entry._2)

    forAll(gen) {
      case (entries, k, v) =>
        val lru =
          buildLRU(ImmutableLRU[String, Double](entries.size + 1), entries)

        // make sure it's there
        assert(lru.get(k)._1 == Some(v))

        // remove once
        val (key, newLru) = lru.remove(k)
        assert(key == Some(v))

        // remove it again
        val (key2, _) = newLru.remove(k)
        assert(key2 == None)

        // still should be able to get it out of the original
        val (key3, _) = lru.remove(k)
        assert(key3 == Some(v))
    }
  }
}

object LRUEntriesGenerator {

  /**
    * Generate a nonempty list of map entries keyed with strings
    * and having any arbitrarily typed values.
    */
  def apply[V: Arbitrary]: Gen[List[(String, V)]] = {
    val gen = for {
      size <- Gen.choose(1, 200)
      uniqueKeys <- Gen.listOfN(size, Gen.identifier).map(_.toSet)
      vals <- Gen.listOfN(uniqueKeys.size, arbitrary[V])
    } yield uniqueKeys.toList.zip(vals)
    gen suchThat (_.nonEmpty)
  }
}