primitive recursive function

PrimRec.hs

{-# LANGUAGE DataKinds, GADTs, EmptyCase, ViewPatterns, TypeFamilies, TypeOperators #-}
{-# LANGUAGE MultiParamTypeClasses, FlexibleInstances #-}

import Control.Monad
import Test.QuickCheck
import Debug.Trace

data Nat = Z | S Nat deriving (Show)
type N0 = Z
type N1 = S N0
type N2 = S N1
type N3 = S N2

type family (+) a b where
  Z + m = m
  (S n) + m = S (n + m)

nat :: Int -> Nat
nat 0 = Z
nat n = S (nat $ n-1)

fromNat :: Nat -> Int
fromNat Z = 0
fromNat (S n) = fromNat n + 1

data Tuple n a where
  TEmpty :: Tuple N0 a
  TSuc :: a -> Tuple n a -> Tuple (S n) a

instance Functor (Tuple n) where
  fmap f TEmpty = TEmpty
  fmap f (TSuc x xs) = TSuc (f x) (fmap f xs)

instance Show a => Show (Tuple n a) where
  show TEmpty = "TEmpty"
  show (TSuc x xs) = show x ++ " %: " ++ show xs

lengthT :: Tuple n a -> Nat
lengthT TEmpty = Z
lengthT (TSuc _ t) = S (lengthT t)

infixr 4 %:
(%:) = TSuc

-- datatype of `p.r. function :: N^p -> N`
data PrimRec p where
  Zero :: PrimRec N0
  Suc :: PrimRec N1
  Proj :: Nat -> PrimRec n
  Comp :: PrimRec m -> Tuple m (PrimRec n) -> PrimRec n
  Induct :: PrimRec n -> PrimRec (S (S n)) -> PrimRec (S n)

class VarLenTuple n where
  varLenTuple :: Tuple n Nat

instance VarLenTuple Z where
  varLenTuple = TEmpty

instance (VarLenTuple n) => VarLenTuple (S n) where
  varLenTuple = Z %: varLenTuple

dimDom :: (VarLenTuple n) => PrimRec n -> Nat
dimDom = lengthT . go where
  go :: (VarLenTuple n) => PrimRec n -> Tuple n Nat
  go _ = varLenTuple

class VarTuple n where
  varTuple' :: Nat -> Tuple n (PrimRec m)

instance VarTuple Z where
  varTuple' _ = TEmpty

instance (VarTuple n) => VarTuple (S n) where
  varTuple' (S n) = attach (Proj n) $ varTuple' n

attach :: a -> Tuple n a -> Tuple (S n) a
attach a (TSuc x xs) = TSuc x (attach a xs)
attach a TEmpty = TSuc a TEmpty

varTuple :: (VarLenTuple n, VarTuple n) => PrimRec n -> Tuple n (PrimRec m)
varTuple = varTuple' . dimDom where

mapProj :: (Nat -> Nat) -> PrimRec n -> PrimRec n
mapProj k Zero = Zero
mapProj k Suc = Suc
mapProj k (Proj i) = Proj $ k i
mapProj k (Comp f fs) = Comp (mapProj k f) (fmap (mapProj k) fs)
mapProj k (Induct f g) = Induct (mapProj k f) (mapProj k g)

infixl 6 %.
(%.) = Comp

(!) :: Show a => Tuple n a -> Nat -> a
(TSuc n _) ! Z = n
(TSuc _ t) ! (S n) = t ! n
k ! a = error $ show k ++ " ! " ++ show a

evalPR :: PrimRec n -> Tuple n Nat -> Nat
evalPR Zero _ = Z
evalPR Suc (TSuc n TEmpty) = S n
evalPR (Proj k) t = t ! k
evalPR (Comp f fs) xs = evalPR f (go fs xs)
 where
   go :: Tuple m (PrimRec n) -> Tuple n Nat -> Tuple m Nat
   go (TSuc f fs) xs = TSuc (evalPR f xs) (go fs xs)
   go TEmpty xs = TEmpty
evalPR (Induct f h) (TSuc Z xs) = evalPR f xs
evalPR g@(Induct f h) (TSuc (S y) xs) = evalPR h (evalPR g (y %: xs) %: y %: xs)
evalPR _ z = case z of {}

runPR :: PrimRec n -> Tuple n Nat -> Int
runPR f ts = fromNat $ evalPR f ts

{-
  g = Comp f fs:
    g(xs) = f(f1(xs), ... ,fs(xs))

  g = Induct f h:
    g(0,xs) = f(xs)
    g(y+1,xs) = h(g(y,xs),y,xs)
-}

infixr 8 %..
(%..) f x = f %. (x %: TEmpty)

extend :: PrimRec N0 -> PrimRec n
extend n = n %. TEmpty

arg :: Int -> PrimRec n
arg n = Proj (nat n)

-- one = suc(zero)
one :: PrimRec N0
one = Suc %.. Zero

-- two = suc(one)
two :: PrimRec N0
two = Suc %.. one

-- add(0,x) = x; add(y+1,x) = suc(add(x,y))
add :: PrimRec N2
add = Induct (arg 0) (Suc %. (arg 0 %: TEmpty))

-- predpr(0) = 0; predpr(y+1) = y
predpr :: PrimRec N1
predpr = Induct Zero (Proj (nat 1))

-- sub(x,y) = sub'(y,x)
-- sub'(x,0) = x; sub(x,y+1) = pred(sub(x,y))
sub :: PrimRec N2
sub = sub' %. (arg 1 %: arg 0 %: TEmpty) where
  sub' = Induct (arg 0) (predpr %. (arg 0 %: TEmpty))

-- mult(0,x) = 0; mult(y+1,x) = add(x,mult(x,y))
mult :: PrimRec N2
mult = Induct (extend Zero) (add %. (arg 2 %: arg 0 %: TEmpty))

-- equalZ(x) = 1 - (1 - x) = sub(1,sub(1,x)) = char[x=0]
equalZ :: PrimRec N1
equalZ = sub %. (extend one %: sub %. (extend one %: arg 0 %: TEmpty) %: TEmpty)

-- equal(x,y) = equalZ((x-y) + (y-x)) = equalZ(add(sub(x-y),sub(y,x))) = char[x=y]
equal :: PrimRec N2
equal = equalZ %.. add %. (sub %. (arg 0 %: arg 1 %: TEmpty) %: sub %. (arg 1 %: arg 0 %: TEmpty) %: TEmpty)

-- leq(x,y) = equalZ(x-y) = char[x<=y]
leq :: PrimRec N2
leq = equalZ %.. sub %. (arg 0 %: arg 1 %: TEmpty)

lneq :: PrimRec N2
lneq = andp %. (leq %. (arg 0 %: arg 1 %: TEmpty) %: notp %.. equal %. (arg 0 %: arg 1 %: TEmpty) %: TEmpty)

-- not(p)(x) = 1 - p(x) = char[not p]
notp :: PrimRec N1
notp = sub %. (extend one %: arg 0 %: TEmpty)

-- and(p,q) = char[p+q=0] = char[p and q]
andp :: PrimRec N2
andp = equalZ %.. add %. (arg 0 %: arg 1 %: TEmpty)

-- or(p,q) = p*q = char[p or q]
orp :: PrimRec N2
orp = mult %. (arg 0 %: arg 1 %: TEmpty)

-- max2(x,y) = x+(y-x)
max2 :: PrimRec N2
max2 = add %. (arg 0 %: sub %. (arg 1 %: arg 0 %: TEmpty) %: TEmpty)

-- min2(x,y) = x+y - max2(x,y)
min2 :: PrimRec N2
min2 = sub %. (add %. (arg 0 %: arg 1 %: TEmpty) %: max2 %. (arg 0 %: arg 1 %: TEmpty) %: TEmpty)


-- prod{r}(0,xs) = 1; prod{r}(z+1,xs) = mult(r(z,xs),prod{r}(z,xs))
prod :: (VarTuple n, VarLenTuple n) => PrimRec (S n) -> PrimRec (S n)
prod r = Induct (extend one) (mult %. (arg 0 %: r' %: TEmpty)) where
  r' = r %. fmap (mapProj S) (varTuple r)

-- existL{r}(y,xs) = char[exist z<y: r(z,xs)] = pi{z<y}(r(z,xs))
existL :: (VarTuple n, VarLenTuple n) => PrimRec (S n) -> PrimRec (S n)
existL r = equalZ %.. prod r

-- sumpr{r}(y) = sum{z<y}(r(z))
-- sumpr{r}(0) = 0; sumpr{r}(y+1) = add(r(y),sumpr{r}(y))
sumpr :: (VarTuple n, VarLenTuple n) => PrimRec (S n) -> PrimRec (S n)
sumpr r = Induct (extend Zero) (add %. (arg 0 %: r' %: TEmpty)) where
  r' = r %. fmap (mapProj S) (varTuple r)

-- forallL{r}(y) = char[forall z<y: r(z)]
-- forallL{r}(0) = char[sum{z<y}(r(z))=0]
forallL :: PrimRec N1 -> PrimRec N1
forallL r = equalZ %.. sumpr r

-- squarep(x) = char[exist z<x+1: z*z=x] = existL{\xz.z*z=x}(x,x)
squarep :: PrimRec N1
squarep = existL r' %. (arg 0 %: arg 0 %: TEmpty) where
  r' = equal %. (mult %. (arg 0 %: arg 0 %: TEmpty) %: arg 1 %: TEmpty)

-- div(x,y) = char[exist z<y+1: x*z=y] = existL{\zxy.x*z=y}(y+1,x,y)
divp :: PrimRec N2
divp = existL r' %. (Suc %.. arg 1 %: arg 0 %: arg 1 %: TEmpty) where
  r' = equal %. (mult %. (arg 0 %: arg 1 %: TEmpty) %: arg 2 %: TEmpty)

-- primep(x) = char[x>1 & not (exist u<x: exist v<x: u*v=x)]
primep :: PrimRec N1
primep = andp %. (lneq %. (extend one %: arg 0 %: TEmpty) %:
                  notp %.. (existL r' %. (arg 0 %: arg 0 %: TEmpty)) %: TEmpty)
  where
    r' :: PrimRec N2 -- (u,x)
    r' = existL r'' %. (arg 1 %: arg 0 %: arg 1 %: TEmpty)

    r'' :: PrimRec N3 -- (v,u,x)
    r'' = equal %. (mult %. (arg 0 %: arg 1 %: TEmpty) %: arg 2 %: TEmpty)

-- mu{p}(y,xs) = min[z:Nat | z<y & p(z,xs)] or y
-- = sum{v<y} g(v,xs)
-- where g(v,xs) = char[exist z<=v: p(z,xs)] = char[exist z<v+1: p(z,xs)]
mu :: (VarTuple n, VarLenTuple n) => PrimRec (S n) -> PrimRec (S n)
mu p = sumpr g %. (varTuple p) where
  f :: Tuple (S n) (PrimRec m) -> Tuple (S n) (PrimRec m)
  f (TSuc x xs) = TSuc (Suc %.. x) xs

  g :: (VarTuple n, VarLenTuple n) => PrimRec (S n)
  g = existL p %. (f $ varTuple p)

-- 0! = 1
-- (n+1)! = (n+1)*n!
factorial :: PrimRec N1
factorial = Induct (extend one) (mult %. (arg 0 %: Suc %.. arg 1 %: TEmpty))

-- pr(0) = 2
-- pr(x+1) = mu{\x. pr(x) < z & prime(z)}(pr(x)!+2,x)
pr :: PrimRec N1
pr = Induct (extend two) (mu p %. (add %. (extend two %: factorial %.. arg 0 %: TEmpty) %: arg 1 %: TEmpty)) where
  p :: PrimRec N2
  p = andp %. (lneq %. (pr %.. arg 1 %: arg 0 %: TEmpty) %: (primep %.. arg 0) %: TEmpty)

-- quickcheck

instance Arbitrary Nat where
  arbitrary = fmap nat arbitrarySizedNatural

char :: Bool -> Nat
char True = Z
char False = S Z

charInt :: Bool -> Int
charInt = fromNat . char

main = do
  forM_ [0..4] $ \n -> do
    print $ runPR pr (nat n %: TEmpty)
  
  quickCheck $ label "add" $ \n m ->
    runPR add (n %: m %: TEmpty) === fromNat n + fromNat m
  quickCheck $ label "predpr" $ \n ->
    runPR predpr (n %: TEmpty) === if fromNat n == 0 then 0 else fromNat n - 1
  quickCheck $ label "sub" $ \n m ->
    runPR sub (n %: m %: TEmpty) === if fromNat n >= fromNat m then fromNat n - fromNat m else 0
  quickCheck $ label "mult" $ \n m ->
    runPR mult (n %: m %: TEmpty) === fromNat n * fromNat m
  quickCheck $ label "equalZ" $ \n ->
    runPR equalZ (n %: TEmpty) === charInt (fromNat n == 0)
  quickCheck $ label "equal" $ \n m ->
    runPR equal (n %: m %: TEmpty) === charInt (fromNat n == fromNat m)
  quickCheck $ label "leq" $ \n m ->
    runPR leq (n %: m %: TEmpty) === charInt (fromNat n <= fromNat m)
  quickCheck $ label "lneq" $ \n m ->
    runPR lneq (n %: m %: TEmpty) === charInt (fromNat n < fromNat m)
  quickCheck $ label "notp" $ \b ->
    runPR notp (char b %: TEmpty) === charInt (not b)
  quickCheck $ label "andp" $ \p q ->
    runPR andp (char p %: char q %: TEmpty) === charInt (p && q)
  quickCheck $ label "orp" $ \p q ->
    runPR orp (char p %: char q %: TEmpty) === charInt (p || q)
  quickCheck $ label "max2" $ \n m ->
    runPR max2 (n %: m %: TEmpty) === (fromNat n `max` fromNat m)
  quickCheck $ label "min2" $ \n m ->
    runPR min2 (n %: m %: TEmpty) === (fromNat n `min` fromNat m)
  quickCheck $ label "sumpr" $ \n ->
    runPR (sumpr (arg 0)) (n %: TEmpty) === sum [1..fromNat n-1]
  quickCheck $ label "squarep" $ \n ->
    runPR squarep (n %: TEmpty) === charInt (fromNat n `elem` fmap (\x -> x*x) [1..fromNat n-1])
  quickCheck $ label "divp" $ \p q ->
    runPR divp (p %: q %: TEmpty) === if fromNat p == 0 then if fromNat q == 0 then 0 else 1 else charInt (fromNat q `mod` fromNat p == 0)
  quickCheckWith stdArgs {maxSuccess = 30} $ label "primep" $ \n ->
    runPR primep (n %: TEmpty) === charInt (isPrime $ fromNat n)
  quickCheck $ label "factorial" $
    runPR factorial (nat 10 %: TEmpty) === factorial' 10

  where
    isPrime 0 = False
    isPrime 1 = False
    isPrime n = all (\x -> n `mod` x /= 0) [2..n-1]

    factorial' 0 = 1
    factorial' n = factorial' (n-1) * n