llvmpy/llvm_cbuilder/tests/test_isprime.py
2012-11-10 18:56:02 -06:00

125 lines
3 KiB
Python

from llvm.core import *
from llvm.passes import *
from llvm.ee import *
from llvm_cbuilder import *
import llvm_cbuilder.shortnames as C
import unittest, logging
def is_prime(x):
if x <= 2:
return True
if (x % 2) == 0:
return False
for y in range(2, int(1 + x**0.5)):
if (x % y) == 0:
return False
return True
def gen_is_prime(mod):
functype = Type.function(C.int, [C.int])
func = mod.add_function(functype, 'isprime')
cb = CBuilder(func)
arg = cb.args[0]
two = cb.constant(C.int, 2)
true = one = cb.constant(C.int, 1)
false = zero = cb.constant(C.int, 0)
with cb.ifelse( arg <= two ) as ifelse:
with ifelse.then():
cb.ret(true)
with cb.ifelse( (arg % two) == zero ) as ifelse:
with ifelse.then():
cb.ret(false)
idx = cb.var(C.int, 3, name='idx')
with cb.loop() as loop:
with loop.condition() as setcond:
setcond( idx < arg )
with loop.body():
with cb.ifelse( (arg % idx) == zero ) as ifelse:
with ifelse.then():
cb.ret(false)
# increment
idx += two
cb.ret(true)
cb.close()
return func
def gen_is_prime_fast(mod):
functype = Type.function(C.int, [C.int])
func = mod.add_function(functype, 'isprime_fast')
cb = CBuilder(func)
arg = cb.args[0]
two = cb.constant(C.int, 2)
true = one = cb.constant(C.int, 1)
false = zero = cb.constant(C.int, 0)
with cb.ifelse( arg <= two ) as ifelse:
with ifelse.then():
cb.ret(true)
with cb.ifelse( (arg % two) == zero ) as ifelse:
with ifelse.then():
cb.ret(false)
idx = cb.var(C.int, 3, name='idx')
sqrt = cb.get_intrinsic(INTR_SQRT, [C.float])
looplimit = one + sqrt(arg.cast(C.float)).cast(C.int)
with cb.loop() as loop:
with loop.condition() as setcond:
setcond( idx < looplimit )
with loop.body():
with cb.ifelse( (arg % idx) == zero ) as ifelse:
with ifelse.then():
cb.ret(false)
# increment
idx += two
cb.ret(true)
cb.close()
return func
class TestIsPrime(unittest.TestCase):
def test_isprime(self):
mod = Module.new(__name__)
lf_isprime = gen_is_prime(mod)
logging.debug(mod)
mod.verify()
exe = CExecutor(mod)
func = exe.get_ctype_function(lf_isprime, 'bool, int')
for x in range(2, 1000):
msg = "Failed at x = %d" % x
self.assertEqual(func(x), is_prime(x), msg)
def test_isprime_fast(self):
mod = Module.new(__name__)
lf_isprime = gen_is_prime_fast(mod)
logging.debug(mod)
mod.verify()
exe = CExecutor(mod)
func = exe.get_ctype_function(lf_isprime, 'bool, int')
for x in range(2, 1000):
msg = "Failed at x = %d" % x
self.assertEqual(func(x), is_prime(x), msg)
if __name__ == '__main__':
unittest.main()