1 files changed, 30 insertions, 2 deletions

diff --git a/ToxCrypto.hs b/ToxCrypto.hs
index e3bc57f9..87f64c9a 100644
--- a/ToxCrypto.hs
+++ b/ToxCrypto.hs
@@ -53,6 +53,7 @@ import Data.ByteString                  as B
 import qualified Data.ByteString.Base16 as Base16
 import qualified Data.ByteString.Char8  as C8
 import Data.Data
+import Data.Functor.Contravariant
 import Data.Kind
 import Data.Ord
 import Data.Serialize as S
@@ -100,8 +101,25 @@ authAndBytes (Encrypted bs) = (auth,bs')
     (as,bs')   = B.splitAt 16 bs
     Right auth = decode as
 
-data Size a = ConstSize Int
-            | VarSize (a -> Int)
+-- | Info about a type's serialized length. Either the length is known
+-- independently of the value, or the length depends on the value.
+data Size a
+    = VarSize (a -> Int)
+    | ConstSize !Int
+    deriving Typeable
+
+instance Contravariant Size where
+  contramap f sz = case sz of
+      ConstSize n -> ConstSize n
+      VarSize g   -> VarSize (\x -> g (f x))
+
+instance Monoid (Size a) where
+    ConstSize x `mappend` ConstSize y = ConstSize (x + y)
+    VarSize f   `mappend` ConstSize y = VarSize $ \x -> f x + y
+    ConstSize x `mappend` VarSize g   = VarSize $ \y -> x + g y
+    VarSize f   `mappend` VarSize g   = VarSize $ \x -> f x + g x
+    mempty = ConstSize 0
+
 
 class Sized a where size :: Size a
 
@@ -111,6 +129,11 @@ instance Sized a => Serialize (Encrypted a) where
             ConstSize n -> Encrypted <$> getBytes (16 + n) -- 16 extra for Poly1305 mac
     put = putEncrypted
 
+instance Sized a => Sized (Encrypted a) where
+    size = case size of
+        ConstSize n -> ConstSize $ n + 16
+        VarSize f   -> VarSize $ \x -> f x + 16
+
 instance (Sized a, Sized b) => Sized (a,b) where
     size = case (size :: Size a, size :: Size b) of
         (ConstSize a , ConstSize b) -> ConstSize $ a + b
@@ -250,6 +273,11 @@ data Assym a = Assym
     }
  deriving (Functor,Foldable,Traversable)
 
+instance Sized a => Sized (Assym a) where
+    size = case size of
+        ConstSize a -> ConstSize $ a + 24 + 32
+        VarSize f   -> VarSize $ \Assym { assymData = x } -> f x + 24 + 32
+
 -- | Field order: senderKey, then nonce This is the format used by
 -- Ping/Pong/GetNodes/SendNodes.
 --