blog

I started work on my ATS math library recently; though I'm not sure it will be a wholehearted success, there were parts that deserve attention.

Let us first consider the following Haskell function to check primality:

hsIsPrime :: Int -> Bool
hsIsPrime 1 = False
hsIsPrime x = all ((/=0) . (x mod)) [2..m]
    where m = floor (sqrt (fromIntegral x :: Float))

This is actually quite nice, but it is also far from "as fast as C". We can rewrite this in ATS as follows:

#include "share/atspre_staload.hats"
staload "libats/libc/SATS/math.sats"
staload UN = "prelude/SATS/unsafe.sats"
fun sqrt_int(k : intGt(0)) : [ m : nat ] int(m) =
  let
    var pre_bound: int = g0float2int(sqrt_float(g0int2float_int_float(k)))
    var bound: [ m : nat ] int(m) = $UN.cast(pre_bound)
  in
    bound
  end
fn is_prime(k : intGt(0)) : bool =
  case+ k of
    | 1 => false
    | k => 
      begin
        let
          fun loop {n : nat}{m : nat} .<max(0,m-n)>. (i : int(n), bound : int(m)) :<> bool =
            if i < bound then
              if k % i = 0 then
                false
              else
                loop(i + 1, bound)
            else
              if i = bound then
                if k % i = 0 then
                  false
                else
                  true
              else
                true
        in
          loop(2, sqrt_int(k))
        end
      end

As can be seen, this is a bit more involved. Some of this is to be expected as we prove termination, but some of it is simply due to using a lower-level language.

Note first the rather involved termination metric here - we use .<max(0,m-n)>. to ensure that it is well-founded. I learned this trick reading the ATS manual, and I actually encourage prospective learners to do this as well; it can be hard to get your hands on code samples.

Next, note the use of :<> in the type signature of loop. This is how ATS manages effects. You can have :<!ntm> for a function that may not terminate, :<!exn> for a function that could throw an exception, etc. So :<> simply means that the function is pure and total (that is, terminating).

Finally, note the use of an existential type in the declaration var bound: .... Existential types are a feature special to ATS; they are not present in any mainstream languages or even Idris. The type [ m : nat ] int(m) can then be passed as an argument to loop, though loop takes the universally quantified {m : nat} ... bound: int(m). Hence, we have a powerful technique that allows us to take a function on an ordinary int and nonetheless prove termination.

At this point, we can write a small Haskell wrapper. This is slightly dry, but I think it will be informative.

#define ATS_MAINATSFLAG 1
extern
fun is_prime_ats { n : nat | n > 0 } : int(n) -> bool =
  "mac#"
implement is_prime_ats (n) =
  is_prime(n)

This is the necessary code on the ATS side of things; it simply exposes the function in the compiled library. We can compile it to C (placing it in cbits per Haskell packaging convention) with:

 $ mkdir -p cbits
 $ patsopt -dd prime.dats -cc > cbits/prime.c

On the Haskell side we have the following:

module Numeric.NumberTheory where
import           Data.Word
import           Foreign.C
foreign import ccall unsafe is_prime_ats :: CInt -> CBool
convertBool :: CBool -> Bool
convertBool = go . fromIntegral
    where
        go :: Word8 -> Bool
        go 0 = False
        go 1 = True
        go _ = False
isPrime :: Int -> Bool
isPrime = convertBool . is_prime_ats . fromIntegral

Our .cabal file will be a bit simple-minded. We simply add the following to our library stanza:

c-sources:           cbits/primes.c
include-dirs:        /usr/local/lib/ats2-postiats-0.3.8/ccomp/runtime/
                   , /usr/local/lib/ats2-postiats-0.3.8/
ghc-options:         -optc-mtune=native -optc-flto -optc-O3

Finally, we can benchmark both functions alongside one another:

This is exactly what we wanted - a 5x speedup compared to pure Haskell, with even more static guarantees to boot! With C, you can get the same speed (or perhaps even faster), but it always feels like an uncanny compromise. ATS allows us to write the functional code we want (with attendant safety guarantees) and still get the speed we need.