Matrix path search using comonads

matrix-search-comonad.hs

{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE OverloadedLists #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE DeriveFoldable #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
module Lib where

import qualified Data.Vector                   as V
import           Data.Functor.Compose
import           Control.Comonad.Representable.Store
import           Data.Functor.Rep
import           Data.Distributive
import           Control.Monad
import           Control.Comonad
import           Data.Bifunctor
import           Data.Foldable
import           Data.Functor.Identity
import qualified Data.Map                      as M
import           Data.Function                            ( on )
import           GHC.TypeLits
import           Data.Proxy

-- 2D grid type with a type-level grid size
newtype Grid (size :: Nat) a = Grid (Compose V.Vector V.Vector a)
  deriving (Functor, Show, Foldable, Eq)

instance (KnownNat n) => Distributive (Grid n) where
  distribute = distributeRep

instance (KnownNat n) => Representable (Grid n) where
  type Rep (Grid n) = (Int, Int)
  index (Grid (Compose v)) (x, y) = v V.! x V.! y
  tabulate f = Grid . Compose $ V.generate gridSize go
    where
      gridSize = fromIntegral $ natVal (Proxy @n)
      go x = V.generate gridSize (\y -> f (x, y))

-- A Path is a list of integers
type Path = [Int]

-- keyed by largest number in path
type Paths = M.Map Int Path


-- 1 2 3
-- 1 9 4
-- 2 3 3
exampleGrid :: Grid 3 Int
exampleGrid = Grid . Compose $ [[1, 2, 3], [1, 9, 4], [2, 3, 3]]

-- each slot starts with a path of itself
start :: Store (Grid 3) (Int, Paths)
start = nest <$> store (index exampleGrid) (0, 0)
  where nest x = (x, M.singleton x [x])

-- If it's possible for a path from a neighbour to proceed to the focus,
-- then do so and add it to the current cell's paths.
addCurrentToPath :: Int -> Paths -> Paths
addCurrentToPath nextInt opts = tryAddNext <$> opts
 where
  tryAddNext :: Path -> Path
  tryAddNext [] = [nextInt]
  tryAddNext xs@(prev : _) | nextInt > prev = nextInt : xs
                           | otherwise      = xs

-- Select the longest path ending in a given Int, no sense in keeping others
useBest :: Path -> Path -> Path
useBest xs ys | length xs > length ys = xs
              | otherwise             = ys

-- Search one step forward towards the solution by collecting paths from
-- the neighbours of each comonadic context
iterateG
  :: forall n
   . (KnownNat n)
  => Store (Grid n) (Int, Paths)
  -> Store (Grid n) (Int, Paths)
iterateG = extend go
 where
  go :: Store (Grid n) (Int, Paths) -> (Int, Paths)
  go w@(extract  -> (ind, opts)) =
    let neighbours = experiment (neighbourIndices (Proxy @n)) w
        neighbouringOptions =
          M.unionsWith useBest (addCurrentToPath ind . snd <$> neighbours)
        best = M.unionWith useBest neighbouringOptions opts
    in  (ind, best)

-- Get the orthogonal neighbours of coordinates in a grid, filtering out
-- indexes which are off the edge of the grid
neighbourIndices
  :: forall n . (KnownNat n) => Proxy n -> (Int, Int) -> [(Int, Int)]
neighbourIndices _ (x, y) = filter
  (\(x', y') -> x' >= 0 && x' < gridSize && y' >= 0 && y' < gridSize)
  [(x - 1, y), (x + 1, y), (x, y - 1), (x, y + 1)]
  where gridSize = fromIntegral $ natVal (Proxy @n)

unStore :: Store (Grid n) (Int, Paths) -> (Grid n) (Int, Paths)
unStore (StoreT (Identity grid) _) = grid

-- Keep running a function until some property holds between iterations
iterateUntil :: (a -> a -> Bool) -> (a -> a) -> a -> a
iterateUntil pred f a =
  let next = f a in if pred a next then next else iterateUntil pred f next

-- iterate until no progress is made, then we have the solution
solve :: (KnownNat n) => Store (Grid n) (Int, Paths) -> [Int]
solve start = reverse . maximumBy (compare `on` length) $ foldMap
  (toList . snd)
  (unStore solved)
  where solved = iterateUntil ((==) `on` unStore) iterateG start