import qualified Data.Map

type LSystemElement = (Char, [Int])
type LSystem = [LSystemElement]
type LSystemRules = Data.Map.Map Char (LSystemElement -> LSystem)

-- Slow
generate :: LSystemRules -> LSystem -> Int -> LSystem
generate rules axiom steps =
    l !! steps
    where
      l = axiom : [x ++ f x | x <- l]
      f = concatMap $ \elem -> Data.Map.findWithDefault (\x -> [x]) (fst elem) rules elem

-- Faster
generate' :: LSystemRules -> LSystem -> Int -> LSystem
generate' rules axiom steps =
    concatMap (iterate f axiom !!) (ind !! steps)
    where
      ind = [0] : [x ++ map (1+) x | x <- ind]
      f = concatMap $ \elem -> Data.Map.findWithDefault (\x -> [x]) (fst elem) rules elem

-- Fastest
generate'' :: LSystemRules -> LSystem -> Int -> LSystem
generate'' rules axiom steps =
    concatMap (iterate f axiom !!) (ind !! steps)
    where
      ind = [0] : [g x | x <- ind]
          where
            g [] = []
            g [x] = [x, x + 1]
            g xs = xs ++ g (drop (length xs `div` 2) xs)
      f = concatMap $ \elem -> Data.Map.findWithDefault (\x -> [x]) (fst elem) rules elem

-- Fast
generate2 :: LSystemRules -> LSystem -> Int -> LSystem
generate2 rules axiom steps =
    uncurry (++) $ l !! steps
    where
      l = ([], axiom) : (axiom, f axiom) : [(x ++ y, y ++ f y) | (x, y) <- drop 1 l]
      f = concatMap $ \elem -> Data.Map.findWithDefault (\x -> [x]) (fst elem) rules elem

-- Tests
a (_, [x, y]) | y <= 2    = [('A', [x + 2, y + 2])]
              | otherwise = [('B', [2]), ('A', [x - 1, y - 1])]

b (_, [x]) | x <= 2    = [('C', [])]
           | otherwise = [('B', [x - 1])]

rules = Data.Map.fromList [('A', a), ('B', b)]
axiom = [('A', [2, 2])]

propLSys steps =
    (lsys, lsys, lsys) == (lsys', lsys'', lsys2)
    where
      lsys = generate rules axiom steps
      lsys' = generate' rules axiom steps
      lsys'' = generate'' rules axiom steps
      lsys2 = generate2 rules axiom steps