Move required language extensions into pragmas for shootout.
[nofib.git] / shootout / binary-trees / Main.hs
1 {-# LANGUAGE BangPatterns #-}
2 --
3 -- The Computer Language Benchmarks Game
4 -- http://benchmarksgame.alioth.debian.org/
5 --
6 -- Contributed by Don Stewart
7 -- Parallelized by Louis Wasserman
8
9 import System.Environment
10 import Control.Monad
11 import System.Mem
12 import Data.Bits
13 import Text.Printf
14 import GHC.Conc
15
16 --
17 -- an artificially strict tree.
18 --
19 -- normally you would ensure the branches are lazy, but this benchmark
20 -- requires strict allocation.
21 --
22 data Tree = Nil | Node !Int !Tree !Tree
23
24 minN = 4
25
26 io s n t = printf "%s of depth %d\t check: %d\n" s n t
27
28 main = do
29 n <- getArgs >>= readIO . head
30 let maxN = max (minN + 2) n
31 stretchN = maxN + 1
32 -- stretch memory tree
33 let c = {-# SCC "stretch" #-} check (make 0 stretchN)
34 io "stretch tree" stretchN c
35
36 -- allocate a long lived tree
37 let !long = make 0 maxN
38
39 -- allocate, walk, and deallocate many bottom-up binary trees
40 let vs = depth minN maxN
41 mapM_ (\((m,d,i)) -> io (show m ++ "\t trees") d i) vs
42
43 -- confirm the the long-lived binary tree still exists
44 io "long lived tree" maxN (check long)
45
46 -- generate many trees
47 depth :: Int -> Int -> [(Int,Int,Int)]
48 depth d m
49 | d <= m = let
50 s = sumT d n 0
51 rest = depth (d+2) m
52 in s `par` ((2*n,d,s) : rest)
53 | otherwise = []
54 where n = bit (m - d + minN)
55
56 -- allocate and check lots of trees
57 sumT :: Int -> Int -> Int -> Int
58 sumT d 0 t = t
59 sumT d i t = a `par` b `par` sumT d (i-1) ans
60 where a = check (make i d)
61 b = check (make (-i) d)
62 ans = a + b + t
63
64 check = check' True 0
65
66 -- traverse the tree, counting up the nodes
67 check' :: Bool -> Int -> Tree -> Int
68 check' !b !z Nil = z
69 check' b z (Node i l r) = check' (not b) (check' b (if b then z+i else z-i) l) r
70
71 -- build a tree
72 make :: Int -> Int -> Tree
73 make i 0 = Node i Nil Nil
74 make i d = Node i (make (i2-1) d2) (make i2 d2)
75 where i2 = 2*i; d2 = d-1