3 files changed, 37 insertions, 24 deletions

diff --git a/dht-client.cabal b/dht-client.cabal
index a6b06358..9a121f23 100644
--- a/dht-client.cabal
+++ b/dht-client.cabal
@@ -90,6 +90,7 @@ library
                     System.Global6
                     Data.PacketQueue
                     Data.Word64Map
+                    Data.Word64RangeMap
                     OnionRouter
                     Network.Tox
                     Network.Tox.Transport
@@ -190,6 +191,7 @@ library
                     , resourcet
                     , blaze-builder
                     , hinotify
+                    , reference
 
   if impl(ghc < 8)
       Build-depends:  transformers
@@ -223,9 +225,9 @@ library
   ghc-prof-options:
 
   if flag(cryptonite-backport)
-          cpp-options: -DCRYPTONITE_BACKPORT
-       hs-source-dirs: cryptonite-backport
-            C-sources: cbits/cryptonite_xsalsa.c, cbits/cryptonite_salsa.c
+        cpp-options: -DCRYPTONITE_BACKPORT
+        hs-source-dirs: cryptonite-backport
+        C-sources: cbits/cryptonite_xsalsa.c, cbits/cryptonite_salsa.c
         other-modules: Crypto.Cipher.Salsa
                        Crypto.Cipher.XSalsa
                        Crypto.ECC.Class
diff --git a/src/Data/Word64RangeMap.hs b/src/Data/Word64RangeMap.hs
index b5cbe9d1..c754c226 100644
--- a/src/Data/Word64RangeMap.hs
+++ b/src/Data/Word64RangeMap.hs
@@ -15,6 +15,8 @@ import Data.Array.MArray (MArray(..))
 import qualified Data.Array.Base as Base
 import Data.Reference
 import Debug.Trace
+import Control.Concurrent.STM
+import Control.Concurrent.STM.TArray
 import Data.Array.IO
 import Data.IORef
 
@@ -58,6 +60,11 @@ instance (Reference r m, MArray ma e m) => MArray (RefArray r ma) e m where
     unsafeWrite (RefArray r) i e = modifyRef r (\x -> Base.unsafeWrite x i e >> return (x,()))
 -}
 
+emptySTMRangeMap :: STM (RangeMap TArray Word8 TVar)
+emptySTMRangeMap = RefArray <$>
+                     (newTVar =<<
+                        newListArray (0,-1) [])
+
 -- | a sample RangeMap for easier debugging
 getX :: IO (RangeMap IOArray Word8 IORef)
 getX = RefArray <$>
diff --git a/src/Network/Tox/Crypto/Handlers.hs b/src/Network/Tox/Crypto/Handlers.hs
index 59e78213..f4b79272 100644
--- a/src/Network/Tox/Crypto/Handlers.hs
+++ b/src/Network/Tox/Crypto/Handlers.hs
@@ -25,6 +25,7 @@ import qualified Data.PacketQueue as PQ
 import Data.Serialize             as S
 import Data.Word
 import qualified Data.Word64Map as W64
+import Data.Word64RangeMap
 import qualified Data.Set         as Set
 import qualified Data.Array.Unboxed as A
 import SensibleDir
@@ -48,12 +49,12 @@ data NetCryptoSessionStatus = Unaccepted | Accepted | Confirmed
 
 type IOHook addr x = addr -> x -> IO (Maybe (x -> x))
 type NetCryptoHook = IOHook NetCryptoSession CryptoData
-type MsgTypeArray = A.UArray Word8 Word16
-type MsgOutMap = W64.Word64Map Word8
+type MsgTypeArray = A.UArray Word8 Word64
+-- type MsgOutMap = RangeMap STArray Word8 STRef
+-- type MsgOutMap = W64.Word64Map Word8
 -- type MsgOutMap = A.UArray Word64 Word8 -- if above is too slow, switch to this, but use reasonable bounds
-
-msgOutMapLookup :: Word64 -> MsgOutMap -> Maybe Word8
-msgOutMapLookup = W64.lookup
+-- msgOutMapLookup :: Word64 -> MsgOutMap -> STM (Maybe Word8)
+-- msgOutMapLookup k mp = return (W64.lookup k mp)
 
 -- | Information, that may be made visible in multiple sessions, as well
 -- as displayed in some way to the user via mutiple views.
@@ -101,7 +102,7 @@ data NetCryptoSession = NCrypto
                                                 -- may not be in ncHooks yet, but they should appear
                                                 -- here if ncUnrecognizedHook will add them to ncHooks
                                                 -- on an as-need basis.
-    , ncOutgoingIdMap      :: TVar MsgOutMap
+    , ncOutgoingIdMap      :: RangeMap TArray Word8 TVar
     , ncAllSessions        :: NetCryptoSessions -- ^ may be needed if one net-crypto session
                                                 -- needs to possibly start another, as is
                                                 -- the case in group chats
@@ -110,7 +111,7 @@ data NetCryptoSession = NCrypto
     , ncBufferStart        :: TVar Word32
     , ncDequeueThread      :: Maybe ThreadId
     , ncPingMachine        :: Maybe PingMachine
-    , ncOutgoingQueue      :: PQ.PacketOutQueue (State,Nonce24,TVar MsgOutMap)
+    , ncOutgoingQueue      :: PQ.PacketOutQueue (State,Nonce24,RangeMap TArray Word8 TVar)
                                                 CryptoMessage
                                                 (CryptoPacket Encrypted)
                                                 CryptoData
@@ -187,7 +188,7 @@ newSessionsState crypto unrechook hooks = do
                             , svConfigDir   = configdir
                             , svDownloadDir = svDownloadDir0
                             }
-                      , msgTypeArray            = allMsgTypes -- todo make this a parameter
+                      , msgTypeArray            = allMsgTypes id -- todo make this a parameter
                       , inboundQueueCapacity    = 200
                       , outboundQueueCapacity   = 400
                       , nextSessionId           = nextSessionId0
@@ -235,7 +236,7 @@ newHandShakeData crypto basenonce (HParam {hpOtherCookie,hpMySecretKey,hpCookieR
 
 type XMessage = CryptoMessage -- todo
 
-ncToWire :: STM (State,Nonce24,TVar MsgOutMap)
+ncToWire :: STM (State,Nonce24,RangeMap TArray Word8 TVar)
          -> Word32{- packet number we expect to recieve -}
          -> Word32{- buffer_end -}
          -> Word32{- packet number -}
@@ -250,8 +251,9 @@ ncToWire getState seqno bufend pktno msg = do
                     GrpMsg KnownLossy _ -> Lossy
                     GrpMsg KnownLossless _ -> Lossless
     (state,n24,msgOutMapVar) <- getState
-    msgOutMap <- readTVar msgOutMapVar
-    case msgOutMapLookup typ64 msgOutMap of
+    -- msgOutMap <- readTVar msgOutMapVar
+    result1 <- lookupInRangeMap typ64 msgOutMapVar
+    case result1 of -- msgOutMapLookup typ64 msgOutMap of
         Just outid -> do
             let setMessageId (OneByte _)   mid = OneByte (toEnum8 mid)
                 setMessageId (TwoByte _ x) mid = TwoByte (toEnum8 mid) x
@@ -323,7 +325,7 @@ freshCryptoSession  sessions
     ncUnrecognizedHook0 <- atomically $ newTVar (defaultUnrecognizedHook sessions)
     ncIncomingTypeArray0 <- atomically $ newTVar (msgTypeArray sessions)
     let idMap = foldl (\mp (x,y) -> W64.insert x y mp) W64.empty (zip [0..255] [0..255])
-    ncOutgoingIdMap0 <- atomically $ newTVar idMap
+    ncOutgoingIdMap0 <- atomically $ emptySTMRangeMap -- atomically $ newTVar idMap
     ncView0 <- atomically $ newTVar (sessionView sessions)
     pktq <- atomically $ PQ.new (inboundQueueCapacity sessions) 0
     bufstart <- atomically $ newTVar 0
@@ -516,8 +518,8 @@ cryptoNetHandler sessions addr (NetCrypto (CryptoPacket nonce16 encrypted)) = do
                         Nothing -> return ()
                     msgTypes <- atomically $ readTVar ncIncomingTypeArray
                     let msgTyp = cd ^. messageType
-                        msgTypMapped16 = msgTypes A.! fromEnum8 (msgID cm)
-                        msgTypMapped = fromWord16 $ msgTypMapped16
+                        msgTypMapped64 = msgTypes A.! fromEnum8 (msgID cm)
+                        msgTypMapped = fromWord64 $ msgTypMapped64
                         isLossy (GrpMsg KnownLossy _) = True
                         isLossy (Msg mid) | lossyness mid == Lossy = True
                         isLossy _ = False
@@ -542,9 +544,9 @@ runCryptoHook session@(NCrypto {ncState, ncHooks,ncSessionSecret,ncTheirSessionP
     flip fix (cd,hookmap) $ \lookupAgain (cd,hookmap) -> do
         msgTypes <- atomically $ readTVar ncIncomingTypeArray
         let msgTyp = cd ^. messageType
-            msgTypMapped16 = msgTypes A.! fromEnum8 (msgID cm)
-            msgTypMapped = fromWord16 $ msgTypMapped16
-        if msgTypMapped16 == 0
+            msgTypMapped64 = msgTypes A.! fromEnum8 (msgID cm)
+            msgTypMapped = fromWord64 $ msgTypMapped64
+        if msgTypMapped64 == 0
          then return Nothing
          else
            case Map.lookup msgTypMapped hookmap of
@@ -574,17 +576,19 @@ runCryptoHook session@(NCrypto {ncState, ncHooks,ncSessionSecret,ncTheirSessionP
 --   in the MessageId space if they have such a thing.
 mkMsgTypes :: [MessageType] -> MsgTypeArray
 mkMsgTypes msgs = let zeros = A.listArray (0,255) (replicate 256 0)
-                      in zeros A.// map (\x -> (toIndex x,toWord16 x)) msgs
+                      in zeros A.// map (\x -> (toIndex x,toWord64 x)) msgs
     where
         toIndex (Msg mid) = fromIntegral . fromEnum  $ mid
         toIndex (GrpMsg KnownLossless nam) = 0x63 -- fromEnum MESSAGE_GROUPCHAT
         toIndex (GrpMsg KnownLossy    nam) = 0xC7 -- fromEnum LOSSY_GROUPCHAT
 
 -- | Handle all Tox messages that this code base is aware of.
-allMsgTypes :: MsgTypeArray
-allMsgTypes = A.listArray (minBound,maxBound) (0:knownMsgs)
+--   The first parameter is a function which is applied to get the values
+--   for keys of unknown nature. Could be either 'id' or 'const 0'
+allMsgTypes :: (Word64 -> Word64) -> MsgTypeArray
+allMsgTypes fDefault = A.listArray (minBound,maxBound) (0:knownMsgs)
   where
-    knownMsgs :: [Word16]
+    knownMsgs :: [Word64]
     knownMsgs =
       concat [ map (fromIntegral . fromEnum) [ PacketRequest .. KillPacket ]
              , map (const 0)                 [ 3 .. 15 ]  -- UnspecifiedPacket