Siteswap.hs

module Siteswap where

import Data.Bits
import Data.Maybe
import Control.Monad

newtype Height = Height { height :: Int } deriving (Show, Eq, Ord)

newtype Pattern = Pattern { pattern :: [Height] } deriving (Show, Eq)

newtype State = State { state :: Int } deriving (Show, Eq, Ord)

-- Throw a single ball to height h from state s

throw :: Height -> State -> Maybe State
throw (Height 0) (State s) | (s .&. 1) == 0 = 
  Just $ State $ s `shiftR` 1
throw (Height h) (State s) | h > 0 && ((1 `shift` h) .&. s) == 0 = 
  Just $ State $ (s `shiftR` 1) .|. (1 `shift` (h - 1))
throw _ _ = Nothing
  
-- Run a sequence of throws and return the final state
  
throwMany :: Pattern -> State -> Maybe State
throwMany p s0 = foldM (flip throw) s0 (pattern p)
  
-- To find valid patterns, we need to find cycles in a graph
  
type Graph v e = ([v], [(v, v, e)])
  
cycles :: (Eq v) => Graph v e -> [(v, [e])]
cycles (vs, es) = [ (v1, p ++ [e]) | (v1, v2, p) <- paths (vs, es), not $ null p, (v2', v3, e) <- es, v2 == v2', v3 == v1 ]
 
paths :: (Eq v) => Graph v e -> [(v, v, [e])]
paths (vs, es) = [ (v1, v2, p) | v1 <- vs, (v2, p) <- paths' v1 es [] ] where
  paths' v es visited = (v, []) : 
    [ (v3, (e:p))
    | (v1, v2, e) <- es
    , v1 == v
	, not $ v2 `elem` visited
	, (v3, p) <- paths' v2 es (v1:visited) ]
    
-- The number of 1 bits in a state represents the number of balls in the air
    
countBits :: Int -> Int
countBits 0 = 0
countBits 1 = 1
countBits n = let (d, r) = n `divMod` 2 in countBits d + countBits r

-- Create a graph from a transition function

generateGraph :: (Eq v) => (e -> v -> Maybe v) -> [v] -> [e] -> Graph v e
generateGraph transition vs es =
  let edges = [ (v1, v2, e) 
			  | v1 <- vs 
			  , v2 <- vs
			  , e <- es
			  , transition e v1 == Just v2] in (vs, edges)

-- The state transition graph for a number of balls with a maximum throw height

graph :: Height -> Int -> Graph State Height
graph (Height maxHeight) numBalls = 
  let maxState 	= (1 `shiftL` maxHeight) - 1
      states 	= map State $ filter ((== numBalls) . countBits) [0..maxState] in
  generateGraph throw states (map Height [0..maxHeight])

-- Enumerate valid patterns for a number of balls with a maximum throw height

validPatterns :: Height -> Int -> [(State, [Height])]
validPatterns maxHeight numBalls = cycles $ graph maxHeight numBalls