2 files changed, 29 insertions, 19 deletions

diff --git a/OnionRouter.hs b/OnionRouter.hs
index 330e083a..2391e13c 100644
--- a/OnionRouter.hs
+++ b/OnionRouter.hs
@@ -12,8 +12,10 @@ import Network.Tox.NodeId
 import Network.Tox.Onion.Transport
 
 import Control.Concurrent.STM
+import Control.Concurrent.STM.TArray
 import Control.Monad
 import Crypto.Random
+import Data.Array.MArray
 import Data.Bits
 import qualified Data.ByteString     as B
 import qualified Data.HashMap.Strict as HashMap
@@ -47,7 +49,7 @@ data OnionRouter = OnionRouter
       -- so we can decrypt the response.
       pendingQueries  :: TVar (Word64Map PendingQuery)
       -- | The current 12 routes that may be assigned to outgoing packets.
-    , routeMap        :: TVar (IntMap RouteRecord)
+    , routeMap        :: TArray Int (Maybe RouteRecord)
       -- | A set of nodes used to query for random route nodes. These aren't
       -- used directly in onion routes, they are queried for route nodes that
       -- are nearby randomly selected ids.
@@ -115,6 +117,12 @@ freshRoute birthday r mrec = Just $ RouteRecord
     , routeBirthTime = birthday
     }
 
+modifyArray :: TArray Int r -> (r -> r) -> Int -> STM ()
+modifyArray a f i = do
+    mx <- readArray a i
+    writeArray a i $ f mx
+{-# INLINE modifyArray #-}
+
 gotResponse :: RouteRecord -> RouteRecord
 gotResponse rr = rr
     { responseCount = succ $ responseCount rr
@@ -136,7 +144,7 @@ newOnionRouter perror = do
         drg <- newTVar drg0
         -- forM_ [0..11] $ \n -> writeTChan chan $ BuildRoute (RouteId n)
         pq <- newTVar W64.empty
-        rm <- newTVar IntMap.empty
+        rm <- newArray (0,11) Nothing
         tn <- newTVar IntMap.empty
         ti <- newTVar HashMap.empty
         tc <- newTVar 0
@@ -164,7 +172,7 @@ forkRouteBuilder or getnodes = do
         forever $ do
             let checkRebuild rid want_build stm = flip orElse stm $ do
                     wanted_ver <- readTVar want_build
-                    current_ver <- fmap routeVersion . IntMap.lookup rid <$> readTVar (routeMap or)
+                    current_ver <- fmap routeVersion <$> readArray (routeMap or) rid
                     writeTChan (routeLog or) $ "ONION checkRebuild "++show(current_ver,wanted_ver)
                     check $ maybe True (< wanted_ver) current_ver
                     return $ BuildRoute $ RouteId rid
@@ -335,9 +343,10 @@ handleEvent getnodes or e@(BuildRoute (RouteId rid)) = do
         return $ getr
     now <- getPOSIXTime
     atomically $ maybe (return ()) -- writeTVar (pendingRoutes or IntMap.! rid) True)
-                       (\r -> do modifyTVar' (routeMap or)
-                                             (IntMap.alter (freshRoute now r) rid)
-                                 v <- routeVersion . (IntMap.! rid) <$> readTVar (routeMap or)
+                       (\r -> do modifyArray (routeMap or)
+                                             (freshRoute now r)
+                                             rid
+                                 v <- routeVersion . fromJust <$> readArray (routeMap or) rid
                                  writeTVar (pendingRoutes or IntMap.! rid) v
                                  )
                        mb
@@ -362,14 +371,14 @@ lookupSender or saddr (Nonce8 w8) = do
 
 lookupRoute :: OnionRouter -> NodeInfo -> RouteId -> IO (Maybe OnionRoute)
 lookupRoute or ni (RouteId rid) = do
-    mb <- atomically $ IntMap.lookup rid <$> readTVar (routeMap or)
+    mb <- atomically $ readArray (routeMap or) rid
     return $ storedRoute <$> mb
 
 lookupTimeout :: OnionRouter -> Nonce8 -> OnionDestination r -> STM (OnionDestination RouteId, Int)
 lookupTimeout or n8 (OnionDestination asel ni Nothing) = do
     let RouteId rid = routeId (nodeId ni)
-    mrr <- IntMap.lookup rid <$> readTVar (routeMap or)
-    readTVar (routeMap or) >>= \rm -> writeTChan (routeLog or) $ "ONION lookupTimeout " ++ unwords [show rid,show (IntMap.keys rm)]
+    mrr <- readArray (routeMap or) rid
+    writeTChan (routeLog or) $ unwords ["ONION lookupTimeout " ,show rid]
     case mrr of
         Just rr -> return ( OnionDestination asel ni (Just $ RouteId rid), timeoutForRoute rr)
         Nothing -> return ( OnionDestination asel ni Nothing             , 0 )
@@ -378,7 +387,7 @@ hookQueries :: OnionRouter -> (tid -> Nonce8)
                            -> TransactionMethods d tid (OnionDestination RouteId) x
                            -> TransactionMethods d tid (OnionDestination RouteId) x
 hookQueries or t8 tmethods = TransactionMethods
-    { dispatchRegister = \mvar od d -> do --  :: MVar x -> d -> STM (tid, d)
+    { dispatchRegister = \mvar od d -> {-# SCC "hookQ.dispatchRegister" #-} do --  :: MVar x -> d -> STM (tid, d)
         (tid,d') <- dispatchRegister tmethods mvar od d
         let Nonce8 w8 = t8 tid
             ni = onionNodeInfo od
@@ -386,11 +395,11 @@ hookQueries or t8 tmethods = TransactionMethods
             od' = case od of
                 OnionDestination {} -> od { onionRouteSpec = Just rid }
                 OnionToOwner a b    -> OnionToOwner a b -- Type cast.
-        wanted <- readTVar (pendingRoutes or IntMap.! ridn)
-        mr <- IntMap.lookup ridn <$> readTVar (routeMap or)
+        wanted <- {-# SCC "hookQ.wanted" #-} (readTVar (pendingRoutes or IntMap.! ridn))
+        mr <- {-# SCC "hookQ.mr_action" #-} (readArray (routeMap or) ridn)
         -- Block query until a route is ready.
         check $ fromMaybe False $ do
-            RouteRecord{routeVersion=rv} <- mr
+            RouteRecord{routeVersion=rv} <- {-# SCC "hookQ.mr" #-} mr
             return $ wanted <= rv
         let pq = PendingQuery { pendingDestination = od'
                               , pendingVersion = maybe 0 routeVersion mr
@@ -400,7 +409,7 @@ hookQueries or t8 tmethods = TransactionMethods
         modifyTVar' (pendingQueries or) (W64.insert w8 pq)
         writeTChan (routeLog or) $ "ONION query add " ++ unwords [ show (Just $ pendingVersion pq,w8), ":=", show ni ]
         return (tid,d')
-    , dispatchResponse = \tid x d -> do -- :: tid -> x -> d -> STM (d, IO ())
+    , dispatchResponse = \tid x d -> {-# SCC "hookQ.dispatchResponse" #-} do -- :: tid -> x -> d -> STM (d, IO ())
         let Nonce8 w8 = t8 tid
         mb <- W64.lookup w8 <$> readTVar (pendingQueries or)
         modifyTVar' (pendingQueries or) (W64.delete w8)
@@ -408,10 +417,10 @@ hookQueries or t8 tmethods = TransactionMethods
             let od = pendingDestination pq
                 RouteId rid = fromMaybe (routeId (nodeId (onionNodeInfo od)))
                                 $ onionRouteSpec od
-            modifyTVar' (routeMap or) (IntMap.adjust gotResponse rid)
+            modifyArray (routeMap or) (fmap gotResponse) rid
         writeTChan (routeLog or) $ "ONION query del " ++ show (fmap pendingVersion mb, w8)
         dispatchResponse tmethods tid x d
-    , dispatchCancel = \tid d -> do -- :: tid -> d -> STM d
+    , dispatchCancel = \tid d -> {-# SCC "hookQ.dispatchCancel" #-} do -- :: tid -> d -> STM d
         let Nonce8 w8 = t8 tid
         mb <- W64.lookup w8 <$> readTVar (pendingQueries or)
         modifyTVar' (pendingQueries or) (W64.delete w8)
@@ -419,13 +428,13 @@ hookQueries or t8 tmethods = TransactionMethods
             let od = pendingDestination pq
                 RouteId rid = fromMaybe (routeId (nodeId (onionNodeInfo od)))
                                 $ onionRouteSpec od
-            mrr <- IntMap.lookup rid <$> readTVar (routeMap or)
+            mrr <- readArray (routeMap or) rid
             forM_ mrr $ \rr -> do
             when (routeVersion rr == pendingVersion pq) $ do
             let expireRoute = modifyTVar' (pendingRoutes or IntMap.! rid) expire
                 expire ver | ver <= succ (pendingVersion pq) = succ (pendingVersion pq)
                            | otherwise                       = ver
-            modifyTVar' (routeMap or) (IntMap.adjust gotTimeout rid)
+            modifyArray (routeMap or) (fmap gotTimeout) rid
             case rr of
                 RouteRecord{ responseCount = 0
                            , timeoutCount  = c
diff --git a/examples/dhtd.hs b/examples/dhtd.hs
index 3f2a6a63..8e9e7692 100644
--- a/examples/dhtd.hs
+++ b/examples/dhtd.hs
@@ -27,6 +27,7 @@ import Control.Concurrent.STM.TMChan
 import Control.Exception
 import Control.Monad
 import Control.Monad.IO.Class (liftIO)
+import Data.Array.MArray (getAssocs)
 import Data.Bool
 import Data.Char
 import Data.Conduit as C
@@ -930,7 +931,7 @@ clientSession s@Session{..} sock cnum h = do
         ("onion", s) -> cmd0 $ do
             now <- getPOSIXTime
             join $ atomically $ do
-            rm <- readTVar $ routeMap onionRouter
+            rm <- IntMap.fromList . catMaybes . map (\(i,m) -> fmap (i,) m) <$> getAssocs (routeMap onionRouter)
             ts <- readTVar $ trampolineNodes onionRouter
             tcnt <- readTVar $ trampolineCount onionRouter
             icnt <- HashMap.size <$> readTVar (trampolineIds onionRouter)