Number of SK-combinator calculus terms of a given size

CountSK.hs

{-# OPTIONS_GHC -Wno-type-defaults #-}

module CountSK where

import Memo

import Data.Foldable

data Term = S | K | App Term Term

-- The generating function for SK-combinator calculus terms, given the
-- combinatorial class defined by `size`, is:
--
--           1 - sqrt(1 - 8z)
--    A(z) = ----------------
--                  2z
--
-- The radius of convergence of `A(z)` is:
--
--    rho = 0.125
--
-- See section 3 of `https://arxiv.org/pdf/1612.07682.pdf`.
size :: Term -> Int
size S = 0
size K = 0
size (App f a) = 1 + size f + size a

count :: Int -> Integer
count = memo go
  where
  go :: Int -> Integer
  go 0 = 2
  go n
    | even n    = aux [0 .. div (n - 2) 2]
    | otherwise = aux [0 .. div (n - 3) 2] + count (div (n - 1) 2) ^ 2
    where
    aux :: [Int] -> Integer
    aux = sum . map (\i -> 2 * count i * count (n - 1 - i))

--
--    t$ runhaskell SK.hs
--    0: 2
--    1: 4
--    2: 16
--    3: 80
--    4: 448
--    5: 2688
--    6: 16896
--    7: 109824
--    8: 732160
--    9: 4978688
--    10: 34398208
--
-- https://oeis.org/A025225
main :: IO ()
main = do
  let xs = [0..10]
  for_ xs $ \x -> do
    putStrLn $ show x ++ ": " ++ show (count x)

Memo.hs

module Memo where

-- Memoization using `IORef`
-- https://gist.github.com/pedrominicz/0161fb70fe815559ea2546cb979a3b69

import Data.IORef
import System.IO.Unsafe
import qualified Data.Map as M

memoIO :: (Ord a) => (a -> b) -> IO (a -> IO b)
memoIO f = do
  v <- newIORef M.empty
  return $ \x -> do
    m <- readIORef v
    case M.lookup x m of
      Just r -> return r
      Nothing -> do
        let r = f x
        modifyIORef' v (M.insert x r)
        return r

memo :: (Ord a) => (a -> b) -> (a -> b)
memo f = unsafePerformIO . unsafePerformIO (memoIO f)