diff --git a/src/VM.hs b/src/VM.hs
index 345c3cc..4d63ceb 100644
--- a/src/VM.hs
+++ b/src/VM.hs
@@ -47,7 +47,7 @@ lockUnit :: Game -> Game
 lockUnit game = game {
                      board = newBoard,
                      units = otherUnits,
-                     visitedUnits = Set.empty,
+                     visitedUnits = initialVisitedUnits otherUnits,
                      oldLines = clearedLines,
                      score = Game.score game + newScore
                      } where
@@ -56,6 +56,8 @@ lockUnit game = game {
                          (newBoard, clearedLines) = Board.clearLines mergedBoard
                          size = Set.size $ Unit.members currentUnit
                          newScore = moveScore size clearedLines (Game.oldLines game)
+                         initialVisitedUnits [] = Set.empty
+                         initialVisitedUnits (u:us) = Set.singleton (hash u)
 
 testStep = let unit = Unit (0, 0) (Set.fromList [(2,9)])
                board = Board 5 10 (Set.fromList [(0,8),(1,8),(0,9),(1,9),(4,8),(3,9),(4,9)])