1 files changed, 25 insertions, 33 deletions

diff --git a/src/Network/Tox/NodeId.hs b/src/Network/Tox/NodeId.hs
index 2ad17616..95604108 100644
--- a/src/Network/Tox/NodeId.hs
+++ b/src/Network/Tox/NodeId.hs
@@ -28,6 +28,7 @@ module Network.Tox.NodeId
     , sampleNodeId) where
 
 import Control.Applicative
+import Control.Arrow
 import Control.Monad
 import Crypto.Error.Types                     (CryptoFailable (..),
                                                throwCryptoError)
@@ -87,42 +88,35 @@ packPublicKey ws = BA.allocAndFreeze (8 * length ws) $
         x:xs -> do poke ptr (toBE64 x)
                    loop xs (plusPtr ptr 8)
 
-newtype NodeId = NodeId [Word64]
- deriving (Eq,Ord) -- ByteArrayAccess) -- (Eq,Ord,ByteArrayAccess, Bits, Hashable)
+-- We represent the node id redundantly in two formats.  The [Word64] format is
+-- convenient for short-circuiting xor/distance comparisons.  The PublicKey
+-- format is convenient for encryption.
+data NodeId = NodeId [Word64] !(Maybe PublicKey)
+
+instance Eq NodeId where
+ (NodeId ws _) == (NodeId xs _)
+  =         ws == xs
+
+instance Ord NodeId where
+    compare (NodeId ws _) (NodeId xs _) = compare ws xs
 
 instance Sized NodeId where size = ConstSize 32
 
 key2id :: PublicKey -> NodeId
-key2id = NodeId . unpackPublicKey
+key2id k = NodeId (unpackPublicKey k) (Just k)
 
 bs2id :: ByteString -> NodeId
-bs2id bs = NodeId . unpackPublicKey $ throwCryptoError . publicKey $ bs
-
-id2key :: NodeId -> PublicKey
-id2key (NodeId key) = throwCryptoError . publicKey $ (packPublicKey key :: BA.Bytes)
-
+bs2id bs = uncurry NodeId . (unpackPublicKey &&& Just) $ throwCryptoError . publicKey $ bs
 
-{-
 id2key :: NodeId -> PublicKey
-id2key recipient = case publicKey recipient of
-    CryptoPassed key -> key
-    -- This should never happen because a NodeId is 32 bytes.
-    CryptoFailed e   -> error ("Unexpected pattern fail: "++show e)
+id2key (NodeId ws (Just key)) = key
+id2key (NodeId key Nothing)   = throwCryptoError . publicKey $ (packPublicKey key :: BA.Bytes)
 
-key2id :: PublicKey -> NodeId
-key2id pk = case S.decode (BA.convert pk) of
-                Left _ -> error "key2id"
-                Right nid -> nid
-
--}
-
-{-
-instance Ord NodeId where
-    compare (NodeId a) (NodeId b) = compare (unpackPublicKey a) (unpackPublicKey b)
--}
+zeroKey :: PublicKey
+zeroKey = throwCryptoError $ publicKey $ B.replicate 32 0
 
 zeroID :: NodeId
-zeroID = NodeId $ replicate 4 0 -- throwCryptoError $ publicKey $ B.replicate 32 0
+zeroID = NodeId (replicate 4 0) (Just zeroKey)
 
 -- | Convert to and from a Base64 variant that uses .- instead of +/.
 nmtoken64 :: Bool -> Char -> Char
@@ -147,24 +141,22 @@ instance S.Serialize NodeId where
     put nid = putPublicKey $ id2key nid
 
 instance Hashable NodeId where
-    hashWithSalt salt (NodeId key) = hashWithSalt salt (head key)
-
--- instance FiniteBits NodeId where finiteBitSize _ = 256
+    hashWithSalt salt (NodeId ws _) = hashWithSalt salt (head ws)
 
 testNodeIdBit :: NodeId -> Word -> Bool
-testNodeIdBit (NodeId ws) i
+testNodeIdBit (NodeId ws _) i -- TODO: Optmize: use ByteArray key if it's available.
   | fromIntegral i < 256 -- 256 bits
   , (q, r) <- quotRem (fromIntegral i) 64
   = testBit (ws !! q) (63 - r)
   |     otherwise      = False
 
 xorNodeId :: NodeId -> NodeId -> NodeId
-xorNodeId (NodeId xs) (NodeId ys) = NodeId $ zipWith xor xs ys
+xorNodeId (NodeId xs _) (NodeId ys _) = NodeId (zipWith xor xs ys) Nothing
 
 sampleNodeId :: Applicative m => (Int -> m ByteString) -> NodeId -> (Int,Word8,Word8) -> m NodeId
-sampleNodeId gen (NodeId self) (q,m,b)
+sampleNodeId gen (NodeId self k) (q,m,b)
     | q <= 0    = bs2id <$> gen 32
-    | q >= 32   = pure (NodeId self)
+    | q >= 32   = pure (NodeId self k)
     | let (qw,r) = (q+7) `divMod` 8 -- How many Word64 to prepend?
           bw = shiftL (fromIntegral b) (8*(7-r))
           mw = bw - 1 :: Word64
@@ -172,7 +164,7 @@ sampleNodeId gen (NodeId self) (q,m,b)
           h = xor bw (complement mw .&. t0)
     = flip fmap (gen $ 8 * (4 - (qw-1)) ) $ \bs ->
         let (w:ws) = unpackPublicKey bs
-        in NodeId $ hd ++ (h .|. (w .&. mw)) : ws
+        in NodeId (hd ++ (h .|. (w .&. mw)) : ws) Nothing
 
 data NodeInfo = NodeInfo
   { nodeId   :: NodeId