generating maps

mapgen.hs

import System.Environment (getArgs)

import Control.Monad (guard, forM_)

import System.Random (RandomGen)
import qualified System.Random as Random

import Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map

-- Standard default seed.
seed :: Random.StdGen
seed = Random.mkStdGen 0

-- picks a value at random from a list, returning the value and the remainder of the list.
pick :: RandomGen g => [a] -> g -> (a, [a], g)
pick options g = (chosen, pre ++ post, g')
  where (idx, g') = Random.uniformR (0, length options - 1) g
        (pre, chosen : post) = splitAt idx options

-- generate by repeatedly picking an unassigned cell at random and backtracking over its
-- possible terrains. is passed an existing world and list of unassigned cells.
generate :: (Ord k, RandomGen g) =>
            (k -> v -> Map k v -> Bool) ->
            [k] -> [v] -> g -> [Map k v]
generate ok keys values g = loop Map.empty keys g
  where loop known [] g = [known]
        loop known unknown g =
          let (posn, unknown', g') = pick unknown g in
          do terrain <- values
             -- Check consistency
             guard $ ok posn terrain known
             loop (Map.insert posn terrain known) unknown' g'


-- Terrain worlds, subject to no-mountain-sea-adjacency restriction.
data Terrain = Sea | Plain | Mountain deriving (Show, Eq, Ord, Enum, Bounded)
type Posn = (Int, Int)
type World = Map Posn Terrain

generateTerrain :: RandomGen g => Int -> g -> [World]
generateTerrain size = generate (ok size) [(x,y) | x <- [1..size], y <- [1..size]] terrains

-- Is this assignment allowed?
ok :: Int -> Posn -> Terrain -> World -> Bool
-- -- Uncomment this line to remove all constraints.
-- ok _ _ _ _ = True
ok size _ Plain _ = True
ok size p terrain world =
  and $ do p' <- adjacent size p
           return $ case Map.lookup p' world of
             Just Sea -> (terrain /= Mountain)
             Just Mountain -> (terrain /= Sea)
             _ -> True

adjacent :: Int -> Posn -> [Posn]
adjacent size (x,y) = [(x', y) | x' <- adj x] ++ [(x, y') | y' <- adj y]
  where adj i = filter inRange [i+1, i-1]
        inRange i = 1 <= i && i <= size

enumerate :: (Enum a, Bounded a) => [a]
enumerate = [minBound..maxBound]

terrains :: [Terrain]
terrains = enumerate

defaultSize = 10

main :: IO ()
main = do
  args <- getArgs
  let size = if length args < 1
             then defaultSize
             else read (args !! 0)
  forM_ (generateTerrain size seed) $ \world -> do
         putStrLn $ worldString size world


terrainChar :: Terrain -> Char
terrainChar Sea = '~'
terrainChar Plain = '.'
terrainChar Mountain = '^'

worldString :: Int -> World -> String
worldString size world =
  do row <- [1..size]
     [terrainChar (world Map.! (row,col)) | col <- [1..size]] ++ " "