Zigzag conversion

Zigzag.hs

{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE ViewPatterns #-}
{-# OPTIONS_GHC -Wall -Wextra -Werror -Wno-name-shadowing -Wno-incomplete-uni-patterns #-}

-- An implementation of the zigzag conversion (https://mmhaskell.com/blog/2025/6/9/spatial-reasoning-with-zigzag-patterns) that
-- rearranges the string directly, rather by computing indexes
module Zigzag where

import Data.List
import System.Exit (die)

-- run tests with `ghcid --allow-eval` / `ghciwatch --enable-eval`
--
-- $> examples shouldBe

examples :: Monad m => (String -> String -> m ()) -> m ()
examples shouldBe = do
  zigzag 0 "MONDAYMORNINGHASKELL"
    `shouldBe` ""

  zigzag 1 "MONDAYMORNINGHASKELL"
    `shouldBe` "MONDAYMORNINGHASKELL"

  zigzag 2 "MONDAYMORNINGHASKELL"
    `shouldBe` "M\ \N\ \A\ \M\ \R\ \I\ \G\ \A\ \K\ \L\
               \O\ \D\ \Y\ \O\ \N\ \N\ \H\ \S\ \E\ \L"

  zigzag 3 "MONDAYMORNINGHASKELL"
    `shouldBe` "M\     \A\     \R\     \G\     \K\
               \O\ \D\ \Y\ \O\ \N\ \N\ \H\ \S\ \E\ \L\
               \N\     \M\     \I\     \A\     \L"

  zigzag 4 "MONDAYMORNINGHASKELL"
    `shouldBe` "M\         \M\         \G\         \L\
               \O\     \Y\ \O\     \N\ \H\     \E\ \L\
               \N\ \A\     \R\ \I\     \A\ \K\
               \D\         \N\         \S"

  zigzag 5 "MONDAYMORNINGHASKELL"
    `shouldBe` "M\             \R\              \K\
               \O\         \O\ \N\          \S\ \E\
               \N\     \M\     \I\      \A\     \L\
               \D\ \Y\         \N\ \H\          \L\
               \A\             \G"

shouldBe :: String -> String -> IO ()
shouldBe actual expected
  | actual == expected = putStrLn "✓"
  | otherwise = die do
      "✗ expected " <> show expected <> " but found " <> show actual

zigzag :: Word -> [a] -> [a]
zigzag 0 _ = []
zigzag 1 as = as
zigzag (fromEnum . subtract 1 -> m) as = top <> concat mid <> bot
  where
    ~(top : down, bot : diag) = splitAt m . take (2 * m) $ transpose (chunksOf (2 * m) as) <> repeat []
    mid = zipWith alternate down (reverse diag)

alternate :: [a] -> [a] -> [a]
alternate [] _ = []
alternate (a : as) bs = a : alternate bs as

chunksOf :: Int -> [a] -> [[a]]
chunksOf _ [] = []
chunksOf n as = case splitAt n as of
  ~(xs, ys) -> xs : chunksOf n ys