2 files changed, 75 insertions, 17 deletions

diff --git a/word64-map/src/Data/Word64Map.hs b/word64-map/src/Data/Word64Map.hs
index adc9c27e..4437ca05 100644
--- a/word64-map/src/Data/Word64Map.hs
+++ b/word64-map/src/Data/Word64Map.hs
@@ -1,7 +1,28 @@
 {-# LANGUAGE RankNTypes          #-}
 {-# LANGUAGE ScopedTypeVariables #-}
 {-# LANGUAGE UnboxedTuples       #-}
-module Data.Word64Map where
+{-# LANGUAGE CPP                 #-}
+-- |  This is a wrapper around 'Data.IntMap', from the containers package, but with a
+--    guaranteed bitwidth of 64 bits.
+--
+--    On 32 bit platforms, this is currently accomplished simply by composing two IntMaps.
+--
+--    On obvious 64 bit platforms(platform name as shown by System.Info.arch ends in 64), CPP
+--    is used and it is simply a newtype around IntMap.
+--
+--    If the CPP is not defined, the Word64Map will be a composition of two IntMaps, but it should
+--    work anyway provided Int is either 32 or 64 bits. Nevertheless, for highest guaranteed efficiency,
+--    report your platform so it can be detected and the CPP defined accordingly.
+--
+module Data.Word64Map
+        ( Word64Map
+        , fitsInInt
+        , Data.Word64Map.lookup
+        , insert
+        , delete
+        , size
+        , empty
+        ) where
 
 import Data.Bits
 import qualified Data.IntMap as IntMap
@@ -10,20 +31,44 @@ import Data.Monoid
 import Data.Typeable
 import Data.Word
 
--- | Since 'Int' may be 32 or 64 bits, this function is provided as a
+-- | Returns 'True' if the proxied type can be losslessly converted to 'Int' using
+-- 'fromIntegral'.
+--
+-- Since 'Int' may be 32 or 64 bits, this function is provided as a
 -- convenience to test if an integral type, such as 'Data.Word.Word64', can be
--- safely transformed into an 'Int' for use with 'IntMap'.
+-- safely transformed into an 'Int'.
 --
--- Returns 'True' if the proxied type can be losslessly converted to 'Int' using
--- 'fromIntegral'.
+-- It should be optimized away at runtime.
 fitsInInt :: forall proxy word. (Bounded word, Integral word) => proxy word -> Bool
 fitsInInt proxy = (original == casted)
  where
     original = div maxBound 2 :: word
     casted   = fromIntegral (fromIntegral original :: Int) :: word
 
+#if KNOWN64
+newtype Word64Map a = Word64Map (IntMap a)
+#else
 newtype Word64Map a = Word64Map (IntMap (IntMap a))
+#endif
+
 
+#if KNOWN64
+size :: Word64Map a -> Int
+size (Word64Map mp) = IntMap.size mp
+{-# INLINE size #-}
+empty :: Word64Map a
+empty = Word64Map (IntMap.empty)
+{-# INLINE empty #-}
+lookup :: Word64 -> Word64Map b -> Maybe b
+lookup key (Word64Map mp) = IntMap.lookup (fromIntegral key) mp
+{-# INLINE lookup #-}
+insert :: Word64 -> b -> Word64Map b -> Word64Map b
+insert key val (Word64Map mp) = Word64Map (IntMap.insert (fromIntegral key) val mp)
+{-# INLINE insert #-}
+delete :: Word64 -> Word64Map b -> Word64Map b
+delete key (Word64Map m) = Word64Map (IntMap.delete (fromIntegral key) m)
+{-# INLINE delete #-}
+#else
 size :: Word64Map a -> Int
 size (Word64Map m) = getSum $ foldMap (\n -> Sum (IntMap.size n)) m
 
@@ -32,20 +77,20 @@ empty = Word64Map IntMap.empty
 
 -- Warning: This function assumes an 'Int' is either 64 or 32 bits.
 keyFrom64 :: Word64 -> (# Int,Int #)
-keyFrom64 w8 =
+keyFrom64 key =
     if fitsInInt (Proxy :: Proxy Word64)
-        then (# fromIntegral w8              , 0               #)
-        else (# fromIntegral (w8 `shiftR` 32), fromIntegral w8 #)
+        then (# fromIntegral key              , 0               #)
+        else (# fromIntegral (key `shiftR` 32), fromIntegral key #)
 {-# INLINE keyFrom64 #-}
 
 lookup :: Word64 -> Word64Map b -> Maybe b
-lookup w8 (Word64Map m) | (# hi,lo #) <- keyFrom64 w8 = do
+lookup key (Word64Map m) | (# hi,lo #) <- keyFrom64 key = do
     m' <- IntMap.lookup hi m
     IntMap.lookup lo m'
 {-# INLINE lookup #-}
 
 insert :: Word64 -> b -> Word64Map b -> Word64Map b
-insert w8 b (Word64Map m) | (# hi,lo #) <- keyFrom64 w8
+insert key b (Word64Map m) | (# hi,lo #) <- keyFrom64 key
     = Word64Map $ IntMap.alter (Just . maybe (IntMap.singleton lo b)
                                              (IntMap.insert lo b))
                                hi
@@ -53,7 +98,7 @@ insert w8 b (Word64Map m) | (# hi,lo #) <- keyFrom64 w8
 {-# INLINE insert #-}
 
 delete :: Word64 -> Word64Map b -> Word64Map b
-delete w8 (Word64Map m) | (# hi,lo #) <- keyFrom64 w8
+delete key (Word64Map m) | (# hi,lo #) <- keyFrom64 key
     = Word64Map $ IntMap.alter (maybe Nothing
                                       (\m' -> case IntMap.delete lo m' of
                                                 m'' | IntMap.null m'' -> Nothing
@@ -61,6 +106,4 @@ delete w8 (Word64Map m) | (# hi,lo #) <- keyFrom64 w8
                                hi
                                m
 {-# INLINE delete #-}
-
-
-
+#endif
diff --git a/word64-map/word64-map.cabal b/word64-map/word64-map.cabal
index 938d2e22..576046e8 100644
--- a/word64-map/word64-map.cabal
+++ b/word64-map/word64-map.cabal
@@ -3,18 +3,25 @@
 
 name:                word64-map
 version:             0.1.0.0
--- synopsis:
--- description:
+synopsis: 64-bit Version of IntMap
+description: This is a wrapper around 'Data.IntMap', from the containers package, but with a
+             guaranteed bitwidth of 64 bits.
+
+             On 32 bit platforms, this is currently accomplished simply by composing two IntMaps.
+
+             On obvious 64 bit platforms(platform name as shown by System.Info.arch ends in 64) 
+             it is simply a newtype of IntMap.
 license:             BSD3
 license-file:        LICENSE
 author:              James Crayne
-maintainer:          jim.crayne@gmail.com
+maintainer:          James Crayne <jim.crayne@gmail.com>
 -- copyright:
 -- category:
 build-type:          Simple
 extra-source-files:  CHANGELOG.md
 cabal-version:       >=1.10
 
+
 library
   exposed-modules:     Data.Word64Map
   -- other-modules:
@@ -24,3 +31,11 @@ library
         , containers
   hs-source-dirs:      src
   default-language:    Haskell2010
+  if arch(x86_64)
+    Cpp-options: -DKNOWN64
+  if arch(ppc64)
+    Cpp-options: -DKNOWN64
+  if arch(ia64)
+    Cpp-options: -DKNOWN64
+  if arch(arm64)
+    Cpp-options: -DKNOWN64