Add head to Lists

src/List.ts

@@ -1,8 +1,10 @@
 import { Predicate, Transform, Option } from './constants'
 import ArrayNode from './persistentVector/ArrayNode'
+import size from './listHelpers/size'
 import push from './listHelpers/push'
 import pop from './listHelpers/pop'
 import set from './listHelpers/set'
+import iterate from './listHelpers/iterate'
 import LeafNode, { emptyNode } from './persistentVector/LeafNode'
 
 let EMPTY_LIST: List<any>
@@ -19,26 +21,29 @@ export function emptyList<T>(): List<T> {
 export function makeList<T>(
   tail: LeafNode<T>,
   root?: ArrayNode<T>,
+  head?: LeafNode<T>,
   forceCreation = false
 ): List<T> {
   if (
     !forceCreation &&
-    (!root && !tail.size)
+    (!root && (!head || !head.size) && !tail.size)
   ) {
     return emptyList<T>()
   }
 
   const res = Object.create(List.prototype)
   res.tail = tail
   res.root = root
-  res.size = tail.size + (root ? root.size : 0)
+  res.head = head
+  res.size = size<T>(tail, root, head)
   return res
 }
 
 export default class List<T> {
   tail: LeafNode<T>
   size: number
   root?: ArrayNode<T>
+  head?: LeafNode<T>
   owner?: Object
 
   constructor() {
@@ -50,15 +55,20 @@ export default class List<T> {
   }
 
   get(key: number, notSetVal?: T): Option<T> {
-    const { size, root, tail } = this
+    const { size, root, head, tail } = this
 
     if (key < 0 || key >= size) {
       return notSetVal
-    } else if (key >= (root ? root.size : 0)) {
-      return tail.get(key, notSetVal)
+    } else if (head && key < head.size) {
+      return head.get(key, notSetVal)
     }
 
-    return (root as ArrayNode<T>).get(key, notSetVal)
+    const index = key - (head ? head.size : 0)
+    if (root && index < root.size) {
+      return root.get(index, notSetVal)
+    }
+
+    return tail.get(index, notSetVal)
   }
 
   set(key: number, value: T): List<T> {
@@ -76,27 +86,32 @@ export default class List<T> {
   map<G>(
     transform: Transform<number, Option<T>, Option<G>>
   ): List<G> {
-    let rootSize: number
+    let head: Option<LeafNode<G>>
     let root: Option<ArrayNode<G>>
 
+    let offset = 0
+
+    if (this.head) {
+      head = this.head.map<G>(0, transform, this.owner) as LeafNode<G>
+      offset = head.size
+    }
+
     if (this.root) {
       root = this.root.map<G>(0, transform, this.owner) as ArrayNode<G>
-      rootSize = this.root.size
-    } else {
-      root = undefined
-      rootSize = 0
+      offset = offset + root.size
     }
 
-    const tail = this.tail.map<G>(rootSize, transform, this.owner) as LeafNode<G>
+    const tail = this.tail.map<G>(offset, transform, this.owner) as LeafNode<G>
 
     if (this.owner) {
       const res = (this as List<any>)
-      res.root = root
       res.tail = tail
+      res.root = root
+      res.head = head
       return (res as List<G>)
     }
 
-    return makeList<G>(tail, root)
+    return makeList<G>(tail, root, head)
   }
 
   clear(): List<T> {
@@ -108,7 +123,7 @@ export default class List<T> {
       return this
     }
 
-    const res = makeList<T>(this.tail, this.root, true)
+    const res = makeList<T>(this.tail, this.root, this.head, true)
     res.owner = {}
     return res
   }
@@ -132,29 +147,6 @@ export default class List<T> {
     step: Predicate<number, Option<T>>,
     reverse?: boolean
   ) {
-    if (reverse) {
-      const rootSize = this.root ? this.root.size : 0
-      if (this.tail.iterateReverse(rootSize, step) === true) {
-        return true
-      }
-
-      if (this.root && this.root.iterateReverse(0, step) === true) {
-        return true
-      }
-
-      return false
-    }
-
-    let rootSize: number
-    if (this.root) {
-      rootSize = this.root.size
-      if (this.root.iterate(0, step) === true) {
-        return true
-      }
-    } else {
-      rootSize = 0
-    }
-
-    return this.tail.iterate(rootSize, step)
+    return iterate<T>(this, step, reverse)
   }
 }

src/__tests__/__snapshots__/List.test.ts.snap

@@ -1,5 +1,6 @@
 exports[`List pop correctly removes values and shrinks root as necessary 1`] = `
 List {
+  "head": undefined,
   "root": ArrayNode {
     "content": Array [
       LeafNode {
@@ -628,6 +629,7 @@ List {
 
 exports[`List push correctly appends values and expands root as necessary 1`] = `
 List {
+  "head": undefined,
   "root": ArrayNode {
     "content": Array [
       LeafNode {
@@ -1864,6 +1866,7 @@ List {
 
 exports[`List set correctly sets values and expands root as necessary 1`] = `
 List {
+  "head": undefined,
   "owner": undefined,
   "root": ArrayNode {
     "content": Array [

src/listHelpers/iterate.ts

@@ -0,0 +1,51 @@
+import { Predicate, Option } from '../constants'
+import List from '../List'
+
+function iterate<T>(
+  list: List<T>,
+  step: Predicate<number, Option<T>>,
+  reverse?: boolean
+): boolean {
+  let offset: number
+
+  if (reverse) {
+    offset = list.size - list.tail.size
+    if (list.tail.iterateReverse(offset, step)) {
+      return true
+    }
+
+    if (list.root) {
+      offset = offset - list.root.size
+      if (list.root.iterateReverse(offset, step)) {
+        return true
+      }
+    }
+
+    if (list.head && list.head.iterateReverse(0, step)) {
+      return true
+    }
+
+    return false
+  }
+
+  offset = 0
+  if (list.head) {
+    if (list.head.iterate(0, step)) {
+      return true
+    }
+
+    offset = list.head.size
+  }
+
+  if (list.root) {
+    if (list.root.iterate(offset, step) === true) {
+      return true
+    }
+
+    offset = offset + list.root.size
+  }
+
+  return list.tail.iterate(offset, step)
+}
+
+export default iterate

src/listHelpers/pop.ts

@@ -1,16 +1,22 @@
 import List, { emptyList, makeList } from '../List'
+import ArrayNode from '../persistentVector/ArrayNode'
+import size from './size'
 
 function pop<T>(list: List<T>): List<T> {
-  let tail = list.tail
+  let head = list.head
   let root = list.root
+  let tail = list.tail
 
-  if (tail.size === 1) {
-    if (!root) {
+  if (tail.size <= 1) {
+    if (!root && head) {
+      tail = head
+      head = undefined
+    } else if (!root) {
       return emptyList<T>()
     }
 
-    const lastLeafNode = root.lastLeafNode()
-    root = root.popLeafNode(list.owner)
+    const lastLeafNode = (root as ArrayNode<T>).lastLeafNode()
+    root = (root as ArrayNode<T>).popLeafNode(list.owner)
     tail = lastLeafNode
   } else {
     tail = tail.pop(list.owner)
@@ -19,11 +25,12 @@ function pop<T>(list: List<T>): List<T> {
   if (list.owner) {
     list.tail = tail
     list.root = root
-    list.size = tail.size + (root ? root.size : 0)
+    list.head = head
+    list.size = size<T>(tail, root, head)
     return list
   }
 
-  return makeList<T>(tail, root)
+  return makeList<T>(tail, root, head)
 }
 
 export default pop

src/listHelpers/push.ts

@@ -29,7 +29,7 @@ function push<T>(
     return list
   }
 
-  return makeList<T>(tail, root)
+  return makeList<T>(tail, root, list.head)
 }
 
 export default push

src/listHelpers/set.ts

@@ -10,23 +10,31 @@ function set<T>(
     return list
   }
 
-  let tail = list.tail
+  let head = list.head
   let root = list.root
+  let tail = list.tail
 
   if (key < list.size) {
-    if (root && key < root.size) {
-      root = root.set(key, value, list.owner)
+    if (head && key < head.size) {
+      head = head.set(key, value, list.owner)
     } else {
-      tail = tail.set(key, value, list.owner)
+      const index = key - (head ? head.size : 0)
+
+      if (root && index < root.size) {
+        root = root.set(index, value, list.owner)
+      } else {
+        tail = tail.set(index, value, list.owner)
+      }
     }
 
     if (list.owner) {
       list.tail = tail
       list.root = root
+      list.head = head
       return list
     }
 
-    return makeList<T>(tail, root)
+    return makeList<T>(tail, root, head)
   } else if (key === list.size) {
     return push<T>(list, value)
   }

src/listHelpers/size.ts

@@ -0,0 +1,16 @@
+import ArrayNode from '../persistentVector/ArrayNode'
+import LeafNode from '../persistentVector/LeafNode'
+
+function size<T>(
+  tail: LeafNode<T>,
+  root?: ArrayNode<T>,
+  head?: LeafNode<T>
+) {
+  return (
+    tail.size +
+    (root ? root.size : 0) +
+    (head ? head.size : 0)
+  )
+}
+
+export default size