Add internal docs for palette generator 💡

So that I don't forget how it works :)

flake.nix

@@ -30,13 +30,26 @@
             installPhase = "install -D Stylix/Main $out/bin/palette-generator";
           };
 
+          # Internal documentation
+          palette-generator-haddock = pkgs.stdenvNoCC.mkDerivation {
+            name = "palette-generator-haddock";
+            src = ./palette-generator;
+            buildInputs = [ ghc ];
+            buildPhase =
+              "haddock $src/**/*.hs --html --ignore-all-exports --odir $out";
+            dontInstall = true;
+            dontFixup = true;
+          };
+
           palette-generator-app = utils.lib.mkApp {
             drv = palette-generator;
             name = "palette-generator";
           };
 
         in {
-          packages.palette-generator = palette-generator;
+          packages = {
+            inherit palette-generator palette-generator-haddock;
+          };
           apps.palette-generator = palette-generator-app;
         })) // {
           nixosModules.stylix = { pkgs, ... }@args: {

palette-generator/Ai/Evolutionary.hs

@@ -8,20 +8,35 @@ import Data.List ( mapAccumR, sortBy )
 import Data.Ord ( Down(Down), comparing )
 import System.Random ( RandomGen, randomR )
 
+{- |
+Find every possible combination of two values, with the first value
+coming from one list and the second value coming from a different list.
+-}
 cartesianProduct :: [a] -> [b] -> [(a, b)]
 cartesianProduct = liftA2 (,)
 
+{- |
+Find every possible combination of two values, with both values coming
+from the same list. Values are allowed to be paired with themself.
+-}
 cartesianSquare :: [a] -> [(a, a)]
 cartesianSquare as = as `cartesianProduct` as
 
+-- | Chain a function a set number of times.
 repeatCall :: Int -> (a -> a) -> a -> a
 repeatCall n f = (!! n) . iterate f
 
+-- | Pick a random element from a list using a random generator.
 randomFromList :: (RandomGen r) => r -> [a] -> (a, r)
 randomFromList generator list
   = let (index, generator') = randomR (0, length list - 1) generator
      in (list !! index, generator')
 
+{- |
+Map over a list, passing a random generator into the mapped
+function each time it is called. A random generator is returned
+along with the new list.
+-}
 mapWithGen :: (r -> a -> (r, b)) -> (r, [a]) -> (r, [b])
 mapWithGen = uncurry . mapAccumR
 
@@ -32,28 +47,63 @@ unfoldWithGen f size generator =
       (generator'', a) = f generator'
    in (generator'', a:as)
 
+{- |
+A genotype is a value which is generated by the genetic algorithm.
+
+The environment is used to specify the problem for which
+we are trying to find the optimal genotype.
+-}
 class Species environment genotype where
+  -- | Generate a new genotype at random.
   generate :: (RandomGen r) => environment -> r -> (r, genotype)
+
+  -- | Randomly combine two genotypes.
   crossover :: (RandomGen r) => environment -> r -> genotype -> genotype -> (r, genotype)
+
+  -- | Randomly mutate a genotype using the given environment.
   mutate :: (RandomGen r) => environment -> r -> genotype -> (r, genotype)
+
+  -- | Score a genotype. Higher numbers are better.
   fitness :: environment -> genotype -> Double
 
-data EvolutionConfig = EvolutionConfig { populationSize :: Int
-                                       , survivors :: Int
-                                       , mutationProbability :: Double
-                                       , generations :: Int
-                                       }
+-- | Parameters for the genetic algorithm.
+data EvolutionConfig = EvolutionConfig
+  { -- | The number of genotypes processed on each pass.
+    populationSize :: Int,
+    -- | How many genotypes make it through to the next pass.
+    survivors :: Int,
+    -- | The chance of a genotype being randomly changed
+    --   before crossover. Between 0 and 1.
+    mutationProbability :: Double,
+    -- | Number of passes of the algorithm.
+    generations :: Int
+  }
 
-randomMutation ::
-  (RandomGen r, Species e g) =>
-  e -> EvolutionConfig -> r -> g -> (r, g)
+{- |
+Randomly mutate the given genotype, if the mutation probability
+from the 'EvolutionConfig' says yes.
+-}
+randomMutation :: (RandomGen r, Species e g)
+               => e -- ^ Environment
+               -> EvolutionConfig
+               -> r -- ^ Random generator
+               -> g -- ^ Genotype to mutate
+               -> (r, g)
 randomMutation environment config generator chromosome
   = let (r, generator') = randomR (0.0, 1.0) generator
      in if r <= mutationProbability config
         then mutate environment generator' chromosome
         else (generator', chromosome)
 
-naturalSelection :: (Species e g) => e -> EvolutionConfig -> [g] -> [g]
+{- |
+Select the fittest survivors from a population,
+to be moved to the next pass of the algorithm.
+-}
+naturalSelection :: (Species e g)
+                 => e -- ^ Environment
+                 -> EvolutionConfig
+                 -> [g] -- ^ Original population
+                 -> [g] -- ^ Survivors
 naturalSelection environment config
   = map snd
   . take (survivors config)
@@ -62,9 +112,12 @@ naturalSelection environment config
   -- Down reverses the sort order so that the best fitness comes first
   . map (\genotype -> (Down $ fitness environment genotype, genotype))
 
-evolveGeneration ::
-  (RandomGen r, Species e g) =>
-  e -> EvolutionConfig -> (r, [g]) -> (r, [g])
+-- | Run one pass of the genetic algorithm over a given population.
+evolveGeneration :: (RandomGen r, Species e g)
+                 => e -- ^ Environment
+                 -> EvolutionConfig
+                 -> (r, [g]) -- ^ Random generator, original population
+                 -> (r, [g]) -- ^ New random generator, new population
 evolveGeneration environment config (generator, population)
   = second (naturalSelection environment config)
   $ mapWithGen (randomMutation environment config)
@@ -73,13 +126,24 @@ evolveGeneration environment config (generator, population)
           randomCrossover gen = let (pair, gen') = randomFromList gen pairs
                                  in (uncurry $ crossover environment gen') pair
 
-initialGeneration ::
-  (RandomGen r, Species e g) =>
-  e -> EvolutionConfig -> r -> (r, [g])
+{- |
+Create the initial population, to be fed into the first
+pass of the genetic algorithm.
+-}
+initialGeneration :: (RandomGen r, Species e g)
+                  => e -- ^ Environment
+                  -> EvolutionConfig
+                  -> r -- ^ Random generator
+                 -> (r, [g]) -- ^ New random generator, population
 initialGeneration environment config
   = unfoldWithGen (generate environment) (survivors config)
 
-evolve :: (RandomGen r, Species e g) => e -> EvolutionConfig -> r -> (r, g)
+-- | Run the full genetic algorithm.
+evolve :: (RandomGen r, Species e g)
+       => e -- ^ Environment
+       -> EvolutionConfig
+       -> r -- ^ Random generator
+       -> (r, g) -- ^ New random generator, optimal genotype
 evolve environment config generator
   = second head
   $ repeatCall (generations config) (evolveGeneration environment config)

palette-generator/Data/Colour.hs

@@ -1,10 +1,12 @@
 module Data.Colour ( LAB(..), RGB(..), deltaE, lab2rgb, rgb2lab ) where
 
+-- | Lightness A-B
 data LAB a = LAB { lightness :: a
                  , channelA :: a
                  , channelB :: a
                  }
 
+-- | Red, Green, Blue
 data RGB a = RGB { red :: a
                  , green :: a
                  , blue :: a
@@ -29,6 +31,7 @@ deltaE (LAB l1 a1 b1) (LAB l2 a2 b2) =
       i = deltaL^2 + deltaCkcsc^2 + deltaHkhsh^2
    in if i < 0 then 0 else sqrt i
 
+-- | Convert a 'LAB' colour to a 'RGB' colour
 lab2rgb :: (Floating a, Ord a) => LAB a -> RGB a
 lab2rgb (LAB l a bx) =
   let y = (l + 16) / 116
@@ -48,6 +51,7 @@ lab2rgb (LAB l a bx) =
           , blue = max 0 (min 1 b') * 255
           }
 
+-- | Convert a 'RGB' colour to a 'LAB' colour
 rgb2lab :: (Floating a, Ord a) => RGB a -> LAB a
 rgb2lab (RGB r g b) =
   let r' = r / 255

palette-generator/Stylix/Main.hs

@@ -9,10 +9,14 @@ import System.Exit ( die )
 import System.Random ( mkStdGen )
 import Text.JSON ( encode )
 
-selectColours :: (Floating a, Real a) => V.Vector (LAB a) -> V.Vector (LAB a)
+-- | Run the genetic algorithm to generate a palette from the given image.
+selectColours :: (Floating a, Real a)
+              => V.Vector (LAB a) -- ^ Colours of the source image
+              -> V.Vector (LAB a) -- ^ Generated palette
 selectColours image
   = snd $ evolve image (EvolutionConfig 1000 100 0.5 150) (mkStdGen 0)
 
+-- | Convert a 'DynamicImage' to a simple 'V.Vector' of colours.
 unpackImage :: (Num a) => DynamicImage -> V.Vector (RGB a)
 unpackImage image = do
   let image' = convertRGB8 image
@@ -21,7 +25,9 @@ unpackImage image = do
   let (PixelRGB8 r g b) = pixelAt image' x y
   return $ RGB (fromIntegral r) (fromIntegral g) (fromIntegral b)
 
-loadImage :: String -> IO DynamicImage
+-- | Load an image file.
+loadImage :: String -- ^ Path to the file
+          -> IO DynamicImage
 loadImage input = either error id <$> readImage input
 
 mainProcess :: (String, String) -> IO ()

palette-generator/Stylix/Output.hs

@@ -6,12 +6,20 @@ import Data.Word ( Word8 )
 import Text.JSON ( JSObject, toJSObject )
 import Text.Printf ( printf )
 
+-- | Convert any 'RGB' colour to store integers between 0 and 255.
 toWord8 :: (RealFrac a) => RGB a -> RGB Word8
 toWord8 (RGB r g b) = RGB (truncate r) (truncate g) (truncate b)
 
+{- |
+Convert a colour to a hexdecimal string.
+
+>>> toHex (RGB 255 255 255)
+"#ffffff"
+-}
 toHex :: RGB Word8 -> String
 toHex (RGB r g b) = printf "%02x%02x%02x" r g b
 
+-- | Convert a palette to the JSON format expected by Stylix's NixOS modules.
 makeOutputTable :: (RealFrac a) => V.Vector (RGB a) -> JSObject String
 makeOutputTable
   = toJSObject

palette-generator/Stylix/Palette.hs

@@ -9,21 +9,36 @@ import Data.Vector ( (!), (//) )
 import qualified Data.Vector as V
 import System.Random ( RandomGen, randomR )
 
+-- | Extract the primary scale from a pallete.
 primary :: V.Vector a -> V.Vector a
 primary = V.take 8
 
+-- | Extract the accent colours from a palette.
 accent :: V.Vector a -> V.Vector a
 accent = V.drop 8
 
+{- |
+Combine two palettes by taking a colour from the left,
+then the right, then the left, and so on until we have
+taken enough colours for a new palette.
+-}
 alternatingZip :: V.Vector a -> V.Vector a -> V.Vector a
 alternatingZip = V.izipWith (\i a b -> if even i then a else b)
 
-randomFromVector :: (RandomGen r) => r -> V.Vector a -> (a, r)
+-- | Select a random item from a vector.
+randomFromVector :: (RandomGen r)
+                 => r -- ^ Random generator
+                 -> V.Vector a
+                 -> (a, r) -- ^ Chosen item, new random generator
 randomFromVector generator vector
   = let (index, generator') = randomR (0, V.length vector - 1) generator
      in (vector ! index, generator')
 
 instance (Floating a, Real a) => Species (V.Vector (LAB a)) (V.Vector (LAB a)) where
+  {- |
+  Palettes in the initial population are created by randomly
+  sampling 16 colours from the source image.
+  -}
   generate image = generateColour 16
       where generateColour 0 generator = (generator, V.empty)
             generateColour n generator
@@ -32,23 +47,48 @@ instance (Floating a, Real a) => Species (V.Vector (LAB a)) (V.Vector (LAB a)) w
 
   crossover _ generator a b = (generator, alternatingZip a b)
 
+  {- |
+  Mutation is done by replacing a random slot in the palette with
+  a new colour, which is randomly sampled from the source image.
+  -}
   mutate image generator palette
     = let (index, generator') = randomR (0, 15) generator
           (colour, generator'') = randomFromVector generator' image
        in (generator'', palette // [(index, colour)])
 
   fitness _ palette
-    = realToFrac $ accentDifference - min lightScheme darkScheme
-      where accentDifference = minimum $ do
-              a <- accent palette
-              b <- accent palette
-              return $ deltaE a b
-            lightnesses = V.map lightness palette
-            difference a b = abs $ a - b
-            lightnessError primaryScale accentValue
-              = sum (V.zipWith difference primaryScale $ primary lightnesses)
-              + sum (V.map (difference accentValue) $ accent lightnesses)
-            lightScheme
-              = lightnessError (V.fromList [90, 70, 55, 35, 25, 10, 5, 5]) 40
-            darkScheme
-              = lightnessError (V.fromList [10, 30, 45, 65, 75, 90, 95, 95]) 60
+    = realToFrac $
+      accentDifference -
+      -- Either light schemes or dark themes are allowed.
+      -- We try to converge on whichever theme we are closer to.
+      min lightScheme darkScheme
+      where
+        -- The accent colours should be as different as possible.
+        accentDifference = minimum $ do
+          a <- accent palette
+          b <- accent palette
+          return $ deltaE a b
+
+        -- Helpers for the function below.
+        lightnesses = V.map lightness palette
+        difference a b = abs $ a - b
+
+        lightnessError primaryScale accentValue
+          -- The primary scale's lightnesses should match the given pattern.
+          = sum (V.zipWith difference primaryScale $ primary lightnesses)
+          -- The accent colours should all have the given lightness.
+          + sum (V.map (difference accentValue) $ accent lightnesses)
+
+        {-
+        For light themes, the background is bright and the text is dark.
+        The accent colours are slightly darker.
+        -}
+        lightScheme
+          = lightnessError (V.fromList [90, 70, 55, 35, 25, 10, 5, 5]) 40
+
+        {-
+        For dark themes, the background is dark and the text is bright.
+        The accent colours are slightly brighter.
+        -}
+        darkScheme
+          = lightnessError (V.fromList [10, 30, 45, 65, 75, 90, 95, 95]) 60