LRUCache Swift

import Foundation

public class LRUCache<KeyType: Hashable> {
  private let maxSize: Int
  private var cache: [KeyType: Any] = [:]
  private var priority: LinkedList<KeyType> = LinkedList<KeyType>()
  private var key2node: [KeyType: LinkedList<KeyType>.LinkedListNode<KeyType>] = [:]
  
  public init(_ maxSize: Int) {
    self.maxSize = maxSize
  }
  
  public func get(_ key: KeyType) -> Any? {
    guard let val = cache[key] else {
      return nil
    }
    
    remove(key)
    insert(key, val: val)
    
    return val
  }
  
  public func set(_ key: KeyType, val: Any) {
    if cache[key] != nil {
      remove(key)
    } else if priority.count >= maxSize, let keyToRemove = priority.last?.value {
      remove(keyToRemove)
    }
    
    insert(key, val: val)
  }
  
  private func remove(_ key: KeyType) {
    cache.removeValue(forKey: key)
    guard let node = key2node[key] else {
      return
    }
    priority.remove(node: node)
    key2node.removeValue(forKey: key)
  }
  
  private func insert(_ key: KeyType, val: Any) {
    cache[key] = val
    priority.insert(key, atIndex: 0)
    guard let first = priority.first else {
      return
    }
    key2node[key] = first
  }
}

public final class LinkedList<T> {
  
  public class LinkedListNode<T> {
    var value: T
    var next: LinkedListNode?
    weak var previous: LinkedListNode?
    
    public init(value: T) {
      self.value = value
    }
  }
  
  public typealias Node = LinkedListNode<T>
  
  fileprivate var head: Node?
  
  public init() {}
  
  public var isEmpty: Bool {
    return head == nil
  }
  
  public var first: Node? {
    return head
  }
  
  public var last: Node? {
    if var node = head {
      while let next = node.next {
        node = next
      }
      return node
    } else {
      return nil
    }
  }
  
  public var count: Int {
    if var node = head {
      var c = 1
      while let next = node.next {
        node = next
        c += 1
      }
      return c
    } else {
      return 0
    }
  }
  
  public func node(atIndex index: Int) -> Node? {
    if index >= 0 {
      var node = head
      var i = index
      while node != nil {
        if i == 0 { return node }
        i -= 1
        node = node!.next
      }
    }
    return nil
  }
  
  public subscript(index: Int) -> T {
    let node = self.node(atIndex: index)
    assert(node != nil)
    return node!.value
  }
  
  public func append(_ value: T) {
    let newNode = Node(value: value)
    self.append(newNode)
  }
  
  public func append(_ node: Node) {
    let newNode = LinkedListNode(value: node.value)
    if let lastNode = last {
      newNode.previous = lastNode
      lastNode.next = newNode
    } else {
      head = newNode
    }
  }
  
  public func append(_ list: LinkedList) {
    var nodeToCopy = list.head
    while let node = nodeToCopy {
      self.append(node.value)
      nodeToCopy = node.next
    }
  }
  
  private func nodesBeforeAndAfter(index: Int) -> (Node?, Node?) {
    assert(index >= 0)
    
    var i = index
    var next = head
    var prev: Node?
    
    while next != nil && i > 0 {
      i -= 1
      prev = next
      next = next!.next
    }
    assert(i == 0)  // if > 0, then specified index was too large
    return (prev, next)
  }
  
  public func insert(_ value: T, atIndex index: Int) {
    let newNode = Node(value: value)
    self.insert(newNode, atIndex: index)
  }
  
  public func insert(_ node: Node, atIndex index: Int) {
    let (prev, next) = nodesBeforeAndAfter(index: index)
    let newNode = LinkedListNode(value: node.value)
    newNode.previous = prev
    newNode.next = next
    prev?.next = newNode
    next?.previous = newNode
    
    if prev == nil {
      head = newNode
    }
  }
  
  public func insert(_ list: LinkedList, atIndex index: Int) {
    if list.isEmpty { return }
    var (prev, next) = nodesBeforeAndAfter(index: index)
    var nodeToCopy = list.head
    var newNode: Node?
    while let node = nodeToCopy {
      newNode = Node(value: node.value)
      newNode?.previous = prev
      if let previous = prev {
        previous.next = newNode
      } else {
        self.head = newNode
      }
      nodeToCopy = nodeToCopy?.next
      prev = newNode
    }
    prev?.next = next
    next?.previous = prev
  }
  
  public func removeAll() {
    head = nil
  }
  
  @discardableResult public func remove(node: Node) -> T {
    let prev = node.previous
    let next = node.next
    
    if let prev = prev {
      prev.next = next
    } else {
      head = next
    }
    next?.previous = prev
    
    node.previous = nil
    node.next = nil
    return node.value
  }
  
  @discardableResult public func removeLast() -> T {
    assert(!isEmpty)
    return remove(node: last!)
  }
  
  @discardableResult public func remove(atIndex index: Int) -> T {
    let node = self.node(atIndex: index)
    assert(node != nil)
    return remove(node: node!)
  }
}

extension LinkedList: CustomStringConvertible {
  public var description: String {
    var s = "["
    var node = head
    while node != nil {
      s += "\(node!.value)"
      node = node!.next
      if node != nil { s += ", " }
    }
    return s + "]"
  }
}

extension LinkedList {
  public func reverse() {
    var node = head
    while let currentNode = node {
      node = currentNode.next
      swap(&currentNode.next, &currentNode.previous)
      head = currentNode
    }
  }
}

extension LinkedList {
  public func map<U>(transform: (T) -> U) -> LinkedList<U> {
    let result = LinkedList<U>()
    var node = head
    while node != nil {
      result.append(transform(node!.value))
      node = node!.next
    }
    return result
  }
  
  public func filter(predicate: (T) -> Bool) -> LinkedList<T> {
    let result = LinkedList<T>()
    var node = head
    while node != nil {
      if predicate(node!.value) {
        result.append(node!.value)
      }
      node = node!.next
    }
    return result
  }
}

extension LinkedList {
  convenience init(array: Array<T>) {
    self.init()
    
    for element in array {
      self.append(element)
    }
  }
}

extension LinkedList: ExpressibleByArrayLiteral {
  public convenience init(arrayLiteral elements: T...) {
    self.init()
    
    for element in elements {
      self.append(element)
    }
  }
}

let cache = LRUCache<String>(2)

cache.set("a", val: 1)
cache.set("b", val: 2)
cache.get("a") // returns 1
cache.set("c", val: 3) // evicts key "b"
print(cache)