Undo drawing

jot/JotDrawView.h

@@ -42,6 +42,11 @@
  */
 @property (nonatomic, strong) UIColor *strokeColor;
 
+/**
+ *  Undos previous stroke.
+ */
+- (void)undoLastStroke;
+
 /**
  *  Clears all paths from the drawing, giving a blank slate.
  *

jot/JotDrawView.m

@@ -19,7 +19,7 @@ @interface JotDrawView ()
 
 @property (nonatomic, strong) UIImage *cachedImage;
 
-@property (nonatomic, strong) NSMutableArray *pathsArray;
+@property (nonatomic, strong) NSMutableArray *strokeHistoryArray;
 
 @property (nonatomic, strong) JotTouchBezier *bezierPath;
 @property (nonatomic, strong) NSMutableArray *pointsArray;
@@ -40,8 +40,8 @@ - (instancetype)init
 
         _strokeWidth = 10.f;
         _strokeColor = [UIColor blackColor];
-        
-        _pathsArray = [NSMutableArray array];
+		
+        _strokeHistoryArray = [NSMutableArray array];
 
         _constantStrokeWidth = NO;
 
@@ -58,11 +58,35 @@ - (instancetype)init
 
 #pragma mark - Undo
 
-- (void)clearDrawing
+- (void)undoLastStroke
 {
+    if (self.strokeHistoryArray.count == 0) {
+        return;
+    }
+
     self.cachedImage = nil;
+
+    [self.strokeHistoryArray removeLastObject];
 
-    [self.pathsArray removeAllObjects];
+    self.bezierPath = nil;
+    self.pointsCounter = 0;
+    [self.pointsArray removeAllObjects];
+    self.lastVelocity = self.initialVelocity;
+    self.lastWidth = self.strokeWidth;
+
+    [UIView transitionWithView:self duration:0.2f
+                       options:UIViewAnimationOptionTransitionCrossDissolve
+                    animations:^{
+                        [self redrawBitmap];
+                    }
+                    completion:nil];
+}
+
+- (void)clearDrawing
+{
+    self.cachedImage = nil;
+
+    [self.strokeHistoryArray removeAllObjects];
 
     self.bezierPath = nil;
     self.pointsCounter = 0;
@@ -98,6 +122,7 @@ - (void)drawTouchBeganAtPoint:(CGPoint)touchPoint
     self.lastWidth = self.strokeWidth;
     self.pointsCounter = 0;
     [self.pointsArray removeAllObjects];
+    [self.strokeHistoryArray addObject:[NSMutableArray array]];
     [self.pointsArray addObject:[JotTouchPoint withPoint:touchPoint]];
 }
 
@@ -161,15 +186,16 @@ - (void)drawBitmap
     if (self.cachedImage) {
         [self.cachedImage drawAtPoint:CGPointZero];
     }
-
+
+    [[self.strokeHistoryArray lastObject] addObject:self.bezierPath];
     [self.bezierPath jotDrawBezier];
     self.bezierPath = nil;
 
     if (self.pointsArray.count == 1) {
         JotTouchPoint *touchPoint = [self.pointsArray firstObject];
         touchPoint.strokeColor = self.strokeColor;
         touchPoint.strokeWidth = 1.5f * [self strokeWidthForVelocity:1.f];
-        [self.pathsArray addObject:touchPoint];
+        [self.strokeHistoryArray.lastObject addObject:touchPoint];
         [touchPoint.strokeColor setFill];
         [JotTouchBezier jotDrawBezierPoint:[touchPoint CGPointValue]
                                  withWidth:touchPoint.strokeWidth];
@@ -180,10 +206,24 @@ - (void)drawBitmap
     [self setNeedsDisplay];
 }
 
+- (void)redrawBitmap
+{
+    self.cachedImage = nil;
+
+    UIGraphicsBeginImageContextWithOptions(self.bounds.size, NO, [UIScreen mainScreen].scale);
+
+    [self drawAllPaths];
+
+    self.cachedImage = UIGraphicsGetImageFromCurrentImageContext();
+    UIGraphicsEndImageContext();
+    [self setNeedsDisplay];
+}
+
+
 - (void)drawRect:(CGRect)rect
 {
     [self.cachedImage drawInRect:rect];
-
+    
     [self.bezierPath jotDrawBezier];
 }
 
@@ -202,7 +242,6 @@ - (JotTouchBezier *)bezierPath
 {
     if (!_bezierPath) {
         _bezierPath = [JotTouchBezier withColor:self.strokeColor];
-        [self.pathsArray addObject:_bezierPath];
         _bezierPath.constantWidth = self.constantStrokeWidth;
     }
 
@@ -241,7 +280,7 @@ - (UIImage *)drawAllPathsImageWithSize:(CGSize)size backgroundImage:(UIImage *)b
 
 - (void)drawAllPaths
 {
-    for (NSObject *path in self.pathsArray) {
+	for (NSObject *path in self.allPaths) {
         if ([path isKindOfClass:[JotTouchBezier class]]) {
             [(JotTouchBezier *)path jotDrawBezier];
         } else if ([path isKindOfClass:[JotTouchPoint class]]) {
@@ -252,4 +291,15 @@ - (void)drawAllPaths
     }
 }
 
+#pragma mark - Helpers
+
+- (NSArray *)allPaths
+{
+	NSMutableArray *all = [NSMutableArray array];
+	for (NSArray *item in self.strokeHistoryArray) {
+		[all addObjectsFromArray:item];
+	}
+	return all;
+}
+
 @end

jot/JotViewController.h

@@ -160,6 +160,11 @@ typedef NS_ENUM(NSUInteger, JotViewState){
 
 @property (nonatomic, strong, readonly) JotDrawingContainer *drawingContainer;
 
+/**
+ *  Undoes previous stroke
+ */
+- (void)undo;
+
 /**
  *  Clears all paths from the drawing in and sets the text to an empty string, giving a blank slate.
  */

jot/JotViewController.m

@@ -245,6 +245,11 @@ - (void)setDrawingConstantStrokeWidth:(BOOL)drawingConstantStrokeWidth
 
 #pragma mark - Undo
 
+- (void)undo
+{
+    [self.drawView undoLastStroke];
+}
+
 - (void)clearAll
 {
     [self clearDrawing];