diff --git a/AStar.hs b/AStar.hs
index 7f5ba2c..f3cd835 100644
--- a/AStar.hs
+++ b/AStar.hs
@@ -8,15 +8,10 @@ import qualified Data.Map as M
 import Data.List (foldl')
 import Data.Maybe (fromJust)
 
--- A node in the search
-class (Ord a, Ord b, Num b, Bounded b) => SearchNode a b where
-  -- get the next search node and the cost to reach to it from the current node
-  nextNode ::  a -> [(a, b)]
-
 -- A* algorithm: Find a path from initial node to goal node using a heuristic function.
 -- Returns Nothing if no path found. Else returns Just (path cost, path).
-astar :: SearchNode a b => a ->  a -> (a -> a -> b) -> Maybe (b, [a])
-astar initNode goalNode hueristic =
+astar :: (Ord a, Ord b, Num b) => a -> a -> (a -> [(a, b)]) -> (a -> a -> b) -> Maybe (b, [a])
+astar initNode goalNode nextNode hueristic =
   astar' (PQ.singleton (hueristic initNode goalNode) (initNode, 0))
          S.empty (M.singleton initNode 0) M.empty
   where
@@ -42,10 +37,12 @@ astar initNode goalNode hueristic =
 
         -- Find the successors (with their g and h costs) of the node
         -- which have not been seen yet
-        successors = filter (\(s, g, _) ->
-                                not (S.member s seen') &&
-                                  g < M.findWithDefault maxBound s gscore)
-                     $ successorsAndCosts node gcost
+        successors =
+          filter (\(s, g, _) ->
+                    not (S.member s seen') &&
+                      (not (s `M.member` gscore)
+                        || g < (fromJust . M.lookup s $ gscore)))
+         $ successorsAndCosts node gcost
 
         -- Insert the successors in the open set
         pq'' = foldl' (\q (s, g, h) -> PQ.insert (g + h) (s, g) q) pq' successors
diff --git a/KnightsTravails.hs b/KnightsTravails.hs
index 97b091d..565cd20 100644
--- a/KnightsTravails.hs
+++ b/KnightsTravails.hs
@@ -50,25 +50,24 @@ isValidNotation notation =
        head notation `elem` ['a'..'h'],
        last notation `elem` ['1'..'8']]
 
--- Makes Board an instance of SearchNode for astar to work
-instance SearchNode Board Int where
-  -- Finds the next possible board configurations for one knight's move.
-  -- Move cost is one.
-  nextNode board@(Board {..}) =
-    zip
-      (map (\pos -> board { knightPos = pos })
-       . filter isValidMove
-       . map (\(x, y) -> (fst knightPos + x, snd knightPos + y))
-       $ moves)
-      (repeat 1)
-    where
-      moves = [(1,2), (1,-2), (-1,2), (-1,-2), (2,1), (2,-1), (-2,1), (-2,-1)]
-      isValidMove (x, y) =
-        and [x > 0, x < 9, y > 0, y < 9, not $ (x, y) `S.member` blockedSquares]
+-- Finds the next possible board configurations for one knight's move.
+-- Move cost is one.
+nextKnightPos board@(Board {..}) =
+  zip
+    (map (\pos -> board { knightPos = pos })
+     . filter isValidMove
+     . map (\(x, y) -> (fst knightPos + x, snd knightPos + y))
+     $ moves)
+    (repeat 1)
+  where
+    moves = [(1,2), (1,-2), (-1,2), (-1,-2), (2,1), (2,-1), (-2,1), (-2,-1)]
+    isValidMove (x, y) =
+      and [x > 0, x < 9, y > 0, y < 9, not $ (x, y) `S.member` blockedSquares]
 
 knightAstar heuristic blockedSquares start target =
   fmap (second (map knightPos))
-  $ astar (Board start blockedSquares) (Board target blockedSquares) heuristic
+  $ astar (Board start blockedSquares) (Board target blockedSquares)
+          nextKnightPos heuristic
 
 -- Finds a path from a start square to an end square using BFS
 bfsSearch :: S.Set Square -> Square -> Square -> Maybe (Int, [Square])