diff --git a/app/6.hs b/app/6.hs
index 9ce46ec..40b53de 100644
--- a/app/6.hs
+++ b/app/6.hs
@@ -4,7 +4,6 @@ module Main where
 import Data.Functor
 import Data.List
 import Data.Matrix hiding (trace)
-import Debug.Trace
 
 data Dir where
   N,E,W,S :: Dir
@@ -57,6 +56,12 @@ guardPos s = let
              in
                (y+1,x+1,d)
 
+getMarked :: Matrix Pos -> [(Int,Int)]
+getMarked s = zip [1..] (toLists s) >>= \(y,l) -> zip [1..] l >>= \case
+        (x,Marked) -> [(y,x)]
+        (x,Guard _)  -> [(y,x)]
+        _ -> []
+
 move :: Matrix Pos -> (Int, Int, Dir) -> Matrix Pos
 move s (y,x,d) = let (y',x',d') = case d of
                                  N -> (y-1,x,E)
@@ -70,7 +75,6 @@ move s (y,x,d) = let (y',x',d') = case d of
                      Obstacle -> setElem (Guard d') (y,x) s
                      _ -> setElem (Guard d) (y',x') $ setElem Marked (y,x) s
 
-
 traversal :: Matrix Pos -> Matrix Pos
 traversal s = let i@(y,x,d) = guardPos s in
   if | y == 1       && d == N -> s
@@ -96,29 +100,8 @@ isLoop s v = let i@(y,x,d) = guardPos s
      | elem i v -> True
      | otherwise -> isLoop s' (i:v)
 
-
-traversal' :: Matrix Pos -> [(Int,Int)] -> [(Int,Int)]
-traversal' s v = let i@(y,x,d) = guardPos s
-                     inf = case d of
-                                N -> (y-1,x)
-                                E -> (y,x+1)
-                                W -> (y,x-1)
-                                S -> (y+1,x)
-  in
-  if | y == 1       && d == N -> v
-     | y == nrows s && d == S -> v
-     | x == 1       && d == W -> v
-     | x == ncols s && d == E -> v
-     | otherwise ->
-       let t = (setElem Obstacle inf s) in
-         --trace ("t: " ++ show t) $
-       if isLoop t []
-                        then traversal' (move s i) ((y,x):v)
-                        else traversal' (move s i) v
-
-
 solve2 :: Matrix Pos -> Int
-solve2 s = length . (`traversal'` []) $ s
+solve2 s = length . filter (`isLoop` []) . (<&> flip (setElem Obstacle) s) . (\\ [(\(a,b,_) -> (a,b)) $ guardPos s]) . getMarked . traversal $ s
 
 main :: IO ()
 main = readFile "inputs/6" <&> parse >>= \i ->