Number of untyped lambda calculus terms of a given size

Count.hs

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

module Count where

import Memo

import Data.Foldable

-- Based on section 3.1 of the paper Generating Random Lambda Calculus Terms by
-- Jue Wang.
--
-- https://citeseerx.ist.psu.edu/viewdoc/download?doi=10.1.1.95.2624&rep=rep1&type=pdf

data Term
  = Var Int
  | Lam Term
  | App Term Term

-- Note that there are multiple definitions for the size of a lambda term. This
-- is the definition used in Generating Random Lambda Calculus Terms.
size :: Term -> Integer
size (Var _) = 1
size (Lam b) = 1 + size b
size (App f a) = 1 + size f + size a

-- Counts the number of closed untyped lambda calculus terms of size `n` modulo
-- alpha-equivalence.
--
-- http://oeis.org/A135501
count :: Integer -> Integer
count n = expr (n, 0)
  where
  -- Remove memoization for extreme speeds.
  expr :: (Integer, Integer) -> Integer
  expr = memo go
    where
    go :: (Integer, Integer) -> Integer
    go (0, _) = 0
    go (1, f) = f
    go (n, f) = lam n f + app n f

  lam :: Integer -> Integer -> Integer
  lam n f = expr (n - 1, f + 1)

  app :: Integer -> Integer -> Integer
  app n f
    | even n    = aux [1 .. div (n - 2) 2]
    | otherwise = aux [1 .. div (n - 3) 2] + expr (div (n - 1) 2, f) ^ 2
    where
    aux :: [Integer] -> Integer
    aux = sum . map (\i -> 2 * expr (i, f) * expr (n - 1 - i, f))


-- This definition of size used in Paul Tarau's paper A Hiking Trip Through the
-- Orders of Magnitude: Deriving Efficient Generators for Closed Simply-Typed
-- Lambda Terms and Normal Forms.
--
-- https://arxiv.org/pdf/1608.03912.pdf
size' :: Term -> Integer
size' (Var _) = 0
size' (Lam b) = 1 + size b
size' (App f a) = 1 + size f + size a

--
--    λ> xs = [1..10] ++ [50]
--    λ> :{
--       for_ xs $ \x -> do
--         putStrLn $ show x ++ ": " ++ show (count' x)
--       :}
--    1: 1
--    2: 3
--    3: 14
--    4: 82
--    5: 579
--    6: 4741
--    7: 43977
--    8: 454283
--    9: 5159441
--    10: 63782411
--    50: 996657783344523283417055002040148075226700996391558695269946852267
--
-- https://arxiv.org/pdf/1210.2610.pdf
-- https://oeis.org/A220894
count' :: Integer -> Integer
count' n = expr (n, 0)
  where
  expr :: (Integer, Integer) -> Integer
  expr = memo go
    where
    go :: (Integer, Integer) -> Integer
    go (0, f) = f
    go (n, f) = lam n f + app n f

  lam :: Integer -> Integer -> Integer
  lam n f = expr (n - 1, f + 1)

  app :: Integer -> Integer -> Integer
  app n f
    | even n    = aux [0 .. div (n - 1) 2]
    | otherwise = aux [0 .. div (n - 2) 2] + expr (div (n - 1) 2, f) ^ 2
    where
    aux :: [Integer] -> Integer
    aux = sum . map (\i -> 2 * expr (i, f) * expr (n - 1 - i, f))

-- This definition of size used in the paper Random generation of closed
-- simply-typed λ-terms: a synergy between logic programming and Boltzmann
-- samplers.
--
-- Note that there are two definitions for the size of a lambda term in that
-- paper. This is the definition of size used for deriving the generating
-- function (and thus the Boltzmann sampler) for plain lambda terms.
size'' :: Term -> Integer
size'' (Var x) = toInteger x
size'' (Lam b) = 1 + size'' b
size'' (App f a) = 2 + size'' f + size'' a

-- Counts the number of _plain_ lambda terms, that is, it counts both open and
-- closed terms of a given size, unlike `count` and `count'` which only count
-- closed terms.
count'' :: Int -> Integer
count'' = memo go
  where
  go :: Int -> Integer
  go 0 = 1
  go 1 = 2
  go 2 = 4
  go n = 1 + lam n + app n

  lam :: Int -> Integer
  lam n = count'' (n - 1)

  app :: Int -> Integer
  -- This is horrible, `div (n - 3) 2` and `div (n - 2) 2` always divide an odd
  -- number by two. It works tho.
  app n
    | even n    = aux [0 .. div (n - 3) 2] + count'' (div (n - 2) 2) ^ 2
    | otherwise = aux [0 .. div (n - 2) 2]
    where
    aux :: [Int] -> Integer
    aux = sum . map (\i -> 2 * count'' i * count'' (n - 2 - i))

--
--    $ runhaskell Count.hs
--    1: 0
--    2: 1
--    3: 2
--    4: 4
--    5: 13
--    6: 42
--    7: 139
--    8: 506
--    9: 1915
--    10: 7558
--    300: 473477381975190304152771173386219140737139330572727481574786077988437776485374930756297967309535888631641055198100186060372310059690213121901539467263248370729333568234992881170746462932742300417445775187087736611038605424576071564968041698609452989548022519111112309976379200605183913879075157976088494
--
main :: IO ()
main = do
  let xs = [1..10] ++ [300]
  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)