From 1be65d18ad9dd8f1ff172ad7fcfd7b669cb47d5e Mon Sep 17 00:00:00 2001
From: Max Sokolov <i@maxsokolov.net>
Date: Fri, 2 Sep 2016 21:42:34 +0300
Subject: [PATCH] PrototypeHeightStrategy improvements

---
 Sources/HeightStrategy.swift | 15 +++++++++++++--
 1 file changed, 13 insertions(+), 2 deletions(-)

diff --git a/Sources/HeightStrategy.swift b/Sources/HeightStrategy.swift
index 66eec6f..f0f778d 100644
--- a/Sources/HeightStrategy.swift
+++ b/Sources/HeightStrategy.swift
@@ -49,7 +49,14 @@ public class PrototypeHeightStrategy: CellHeightCalculatable {
             return height
         }
 
-        guard let cell = tableView.dequeueReusableCellWithIdentifier(row.reuseIdentifier) else { return 0 }
+        var prototypeCell = prototypes[row.reuseIdentifier]
+        if prototypeCell == nil {
+
+            prototypeCell = tableView.dequeueReusableCellWithIdentifier(row.reuseIdentifier)
+            prototypes[row.reuseIdentifier] = prototypeCell
+        }
+
+        guard let cell = prototypeCell else { return 0 }
 
         cell.bounds = CGRectMake(0, 0, tableView.bounds.size.width, cell.bounds.height)
 
@@ -66,7 +73,11 @@ public class PrototypeHeightStrategy: CellHeightCalculatable {
     }
 
     public func estimatedHeight(row: Row, path: NSIndexPath) -> CGFloat {
-        return UITableViewAutomaticDimension
+
+        if let estimatedHeight = row.estimatedHeight where estimatedHeight > 0 {
+           return estimatedHeight
+        }
+        return height(row, path: path)
     }
 
     public func invalidate() {