length calculations

Fix: #187

index.js

@@ -1,10 +1,12 @@
 class LRUEntry {
-  constructor (value) {
+  constructor (value, size) {
     this.value = value
+    this.size = size
   }
 }
 
 const asInt = n => ~~n
+const naiveLength = () => 1
 
 class LRUCache {
   constructor (options) {
@@ -19,12 +21,22 @@ class LRUCache {
     this.max = options.max
     this.old = new Map()
     this.current = new Map()
+    this.oldSize = 0
+    this.currentSize = 0
+    this.sizeCalculation = options.length || naiveLength
   }
   get size () {
-    return this.current.size + this.old.size
+    return this.oldSize + this.currentSize
   }
   set (key, value) {
-    this.current.set(key, new LRUEntry(value))
+    const entry = new LRUEntry(value, this.sizeCalculation(value))
+    this.currentSize += entry.size
+    this.current.set(key, entry)
+    this.prune()
+  }
+  promote (key, entry) {
+    this.current.set(key, entry)
+    this.currentSize += entry.size
     this.prune()
   }
   get (key) {
@@ -34,34 +46,44 @@ class LRUCache {
     } else {
       const fromOld = this.old.get(key)
       if (fromOld) {
-        this.current.set(key, fromOld)
-        this.prune()
+        this.promote(key, fromOld)
         return fromOld.value
       }
     }
   }
   delete (key) {
-    this.old.delete(key)
-    this.current.delete(key)
+    const fromOld = this.old.get(key)
+    if (fromOld) {
+      this.old.delete(key)
+      this.oldSize -= fromOld.size
+    }
+    const fromCurrent = this.current.get(key)
+    if (fromCurrent) {
+      this.current.delete(key)
+      this.currentSize -= fromCurrent.size
+    }
   }
   has (key, updateRecency) {
     if (this.current.has(key)) {
       return true
     }
     const oldHas = this.old.get(key)
     if (oldHas && updateRecency) {
-      this.current.set(key, oldHas)
-      this.prune()
+      this.promote(key, oldHas)
     }
     return !!oldHas
   }
   reset () {
     this.old = new Map()
+    this.oldSize = 0
     this.current = new Map()
+    this.currentSize = 0
   }
   prune () {
-    if (this.current.size >= this.max) {
+    if (this.currentSize >= this.max) {
+      this.oldSize = this.currentSize
       this.old = this.current
+      this.currentSize = 0
       this.current = new Map()
     }
   }

test/basic.js

@@ -79,6 +79,11 @@ t.test('basic operation', t => {
   t.equal(c.has(10, true), false)
   t.equal(c.current.size, 1)
   t.equal(c.old.size, 10)
+  c.set(true, 'true')
+  t.equal(c.has(true), true)
+  t.equal(c.get(true), 'true')
+  c.delete(true)
+  t.equal(c.has(true), false)
 
   t.end()
 })

test/size-calculation.js

@@ -0,0 +1,36 @@
+const t = require('tap')
+const LRU = require('../')
+
+t.test('store strings, size = length', t => {
+  const c = new LRU({
+    max: 100,
+    length: n => n.length,
+  })
+  const s = 'x'.repeat(10)
+  for (let i = 0; i < 5; i++) {
+    c.set(i, s)
+  }
+  t.equal(c.size, 50)
+  // the big item goes in, but triggers a prune
+  // we don't preemptively prune until we *cross* the max
+  c.set('big', 'x'.repeat(100))
+  t.equal(c.size, 150)
+  t.equal(c.old.size, 6)
+  t.equal(c.oldSize, 150)
+  t.equal(c.current.size, 0)
+  t.equal(c.currentSize, 0)
+  c.set('big', 'y'.repeat(10))
+  t.equal(c.size, 160)
+  t.equal(c.old.size, 6)
+  t.equal(c.oldSize, 150)
+  t.equal(c.current.size, 1)
+  t.equal(c.currentSize, 10)
+  c.delete('big')
+  t.equal(c.size, 50)
+  t.equal(c.old.size, 5)
+  t.equal(c.oldSize, 50)
+  t.equal(c.current.size, 0)
+  t.equal(c.currentSize, 0)
+
+  t.end()
+})