Lazy calculation of list length; update of https://gist.github.com/rampion/6337593

Magnitude.hs

{-# OPTIONS_GHC -Wall -Wextra -Werror #-}
{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE LambdaCase #-}
-- |
-- Types for lazy length computation and comparison
--
--    >>> length (replicate 5 'a')
--    Positive Magnitude { characteristic = S (S Z), mantissa = 0b01 }
--    >>> fromLength (length (replicate 5 'a'))
--    5
--    >>> length (replicate 5 'a') < length (repeat 'b')
--    True
--
--
module Magnitude
  ( Magnitude.length
  , Length(..)
  , toLength
  , fromLength
  , infinity
  , length1
  , Length1(..)
  , unLength1
  , toLength1
  , fromLength1
  , infinity1
  , Magnitude(..)
  , infinityN
  , Unary(..)
  , Bits(..)
  , Bit(..)
  ) where

import Prelude hiding (length)
import Data.Kind
import Data.List.NonEmpty (NonEmpty(..), nonEmpty)

type Length :: Type
data Length = Zero | Positive (Magnitude 'Z)
  deriving stock (Show, Eq, Ord)

fromLength :: Num a => Length -> a
fromLength = \case
  Zero -> 0
  Positive n -> fromLength1 (Length1 n)

toLength :: Integral a => a -> Either a Length
toLength n = case toLength1 n of
  Right (Length1 p) -> Right (Positive p)
  Left 0 -> Right Zero
  _ -> Left n

type Length1 :: Type
newtype Length1 = Length1 (Magnitude 'Z)
  deriving stock (Show, Eq, Ord)

unLength1 :: Length1 -> Magnitude 'Z
unLength1 (Length1 n) = n

fromLength1 :: Num a => Length1 -> a
fromLength1 = fromMagnitude 1 . unLength1 where
  fromMagnitude :: Num a => a -> Magnitude n -> a
  fromMagnitude !lo = \case
    NextCharacteristic p -> fromMagnitude (2*lo) p
    Mantissa r -> lo + fromBits 0 r

  fromBits :: Num a => a -> Bits n -> a
  fromBits !n = \case
    Nil -> n
    b :. bs -> fromBits (2*n + fromBit b) bs

  fromBit :: Num a => Bit -> a
  fromBit = \case
    O -> 0
    I -> 1

toLength1 :: Integral a => a -> Either a Length1
toLength1 = \n -> if n <= 0 then Left n else Right (Length1 (loop n Nil)) where
  loop :: Integral a => a -> Bits n -> Magnitude n
  loop 1 = Mantissa
  loop n =
    let (q, r) = n `quotRem` 2
        b = if r == 0 then O else I
    in
    NextCharacteristic . loop q . (b :.)

-- |
-- The set of integers greater than or equal to 2ⁿ
type Magnitude :: Unary -> Type
data Magnitude n where
  -- | an integer in the range [2ⁿ,2ⁿ⁺¹)
  Mantissa :: Bits n -> Magnitude n
  -- | an integer in the range [2ⁿ⁺¹,∞)
  NextCharacteristic :: Magnitude ('S n) -> Magnitude n

deriving instance Eq (Magnitude n)
deriving instance Ord (Magnitude n)

-- | This Show instance lies about what constructors Magnitude has, but it's way
-- more legible than the derived version
instance Show (Magnitude n) where
  showsPrec _ m
    = showString "Magnitude { characteristic = "
    . shows (characteristic m) 
    . showString ", mantissa = "
    . withMantissa shows m 
    . showString " }"

withMantissa :: (forall m. Bits m -> r) -> Magnitude n -> r
withMantissa f = \case
  Mantissa bs -> f bs
  NextCharacteristic n -> withMantissa f n

characteristic :: Magnitude n -> Unary
characteristic = \case
  Mantissa _ -> Z
  NextCharacteristic n -> S (characteristic n)

infinity :: Length
infinity = Positive infinityN

infinity1 :: Length1
infinity1 = Length1 infinityN

infinityN :: Magnitude n
infinityN = NextCharacteristic infinityN

type Bit :: Type
data Bit where
  O :: Bit
  I :: Bit

deriving instance Eq Bit
deriving instance Ord Bit
deriving instance Show Bit

-- | Set of n bits, highest order first
type Bits :: Unary -> Type
data Bits n where
  (:.) :: Bit -> Bits n -> Bits ('S n)
  Nil :: Bits 'Z

deriving instance Eq (Bits n)
deriving instance Ord (Bits n)

infixr 4 :.

instance Show (Bits n) where
  showsPrec _ = \bs -> showString "0b" . showBits bs where
    showBits :: Bits m -> ShowS
    showBits = \case
      Nil -> id
      b :. bs -> showBit b . showBits bs

    showBit = (:) . \case
      O -> '0'
      I -> '1'

-- | unary encoding of the natural numbers
type Unary :: Type
data Unary where
  Z :: Unary
  S :: Unary -> Unary

deriving instance Eq Unary
deriving instance Ord Unary
deriving instance Show Unary

-- | Lazily compute the length of a possibly-infinite list
length :: [a] -> Length
length = maybe Zero (Positive . unLength1) . fmap length1 . nonEmpty

-- | Lazily compute the length of a possibly-infinite non-empty list
length1 :: NonEmpty a -> Length1
length1 = Length1 . \(_ :| as) -> fromUnary Nil (unaryLength as) where
  -- lazily compute the length of a list in unary
  --
  -- this isn't really necessary (a list is already a unary representation of
  -- its own length), but it makes it clearer what's going on
  unaryLength :: [a] -> Unary
  unaryLength = \case
    [] -> Z
    _ : as -> S (unaryLength as)

  -- lazily convert a unary value x into the magnitude represention of x+1
  -- carrying over lower bits
  fromUnary :: Bits n -> Unary -> Magnitude n
  fromUnary bs Z = Mantissa bs
  fromUnary bs (S m) = NextCharacteristic
    let ~(m', b) = quotRem2 m in fromUnary (b :. bs) m'

  -- compute (m `quot` 2, m `rem` 2) for a unary value m
  --
  -- quotRem2 m = (q,r) ⇒ m = 2*q + r
  quotRem2 :: Unary -> (Unary, Bit)
  quotRem2 Z = (Z, O)
  quotRem2 (S Z) = (Z, I)
  quotRem2 (S (S m)) =
    let ~(m', b) = quotRem2 m in (S m', b)