Retry logic when "wheel" is not good enough

alba.cabal

@@ -11,73 +11,84 @@ build-type:         Simple
 extra-source-files: README.md
 
 flag dump
-  manual: True
+  manual:  True
   default: True
 
+common warnings
+    ghc-options: -Wall -Wunused-packages -Werror
+                 -fno-warn-missing-pattern-synonym-signatures
+                 -fno-warn-missing-signatures
+                 -fno-warn-name-shadowing
+                 -fno-warn-type-defaults
+                 -fno-warn-unused-imports
+                 -fno-warn-unused-matches
+
 library
-  exposed-modules:
-    ALBA
-  hs-source-dirs:   src
+  import:           warnings
+  exposed-modules:    ALBA
+  hs-source-dirs:     src
   build-depends:
-                    base >=4.9 && <5
-                  , QuickCheck
-                  , base16-bytestring
-                  , bytestring
-                  , cereal
-                  , deepseq
-  default-extensions:  StrictData
-  default-language: Haskell2010
-  ghc-options:      -Wall -Wunused-packages
-  extra-libraries: sodium
-  include-dirs: cbits
-  includes: cbits/blake2b.h
-  c-sources: cbits/blake2b.c
+    , base               >=4.9 && <5
+    , base16-bytestring
+    , bytestring
+    , cereal
+    , deepseq
+    , QuickCheck
+
+  default-extensions: StrictData
+  default-language:   Haskell2010
+  extra-libraries:    sodium
+  include-dirs:       cbits
+  includes:           cbits/blake2b.h
+  c-sources:          cbits/blake2b.c
 
   if flag(dump)
     ghc-options: -ddump-simpl -ddump-stg-final -ddump-to-file
 
 executable alba
+  import:           warnings
   hs-source-dirs: app
   main-is:        Main.hs
   build-depends:
-                  alba
-                , QuickCheck
+                      QuickCheck
+                , alba
                 , base
+                , bytestring
+                , directory
+                , filepath
 
   ghc-options:    -threaded -rtsopts
 
 test-suite test
-    default-language: Haskell2010
-    type:             exitcode-stdio-1.0
-    hs-source-dirs:   test
-    main-is:          Main.hs
-    other-modules:    ALBASpec
-    build-depends:
-                    base >= 4.9 && <5
-                  , QuickCheck
-                  , base16-bytestring
-                  , bytestring
-                  , cereal
-                  , hspec
-                  , alba
-                  , quickcheck-classes
-                  , quickcheck-instances
-
-    build-tool-depends:
-       hspec-discover:hspec-discover
+  import:           warnings
+  default-language:   Haskell2010
+  type:               exitcode-stdio-1.0
+  hs-source-dirs:     test
+  main-is:            Main.hs
+  other-modules:      ALBASpec
+  build-depends:
+                  QuickCheck
+                , alba
+                , base                  >=4.9 && <5
+                , bytestring
+                , cereal
+                , hspec
+                , quickcheck-instances
 
-    ghc-options: -rtsopts -threaded
+  build-tool-depends: hspec-discover:hspec-discover
+  ghc-options:        -rtsopts -threaded
 
 benchmark bench
+  import:           warnings
   hs-source-dirs: bench/
   main-is:        ALBABench.hs
   type:           exitcode-stdio-1.0
   build-depends:
-    , aeson
+    , alba
     , base
     , bytestring
     , criterion
-    , alba
+    , deepseq
     , QuickCheck
 
   ghc-options:    -threaded -rtsopts

app/Main.hs

@@ -2,52 +2,74 @@
 {-# LANGUAGE LambdaCase #-}
 {-# LANGUAGE NamedFieldPuns #-}
 
-import ALBA (Params (..), Verification (Verified), genItems, prove, readProof, verify, writeProof)
+import ALBA (Bytes (..), NoProof (..), Params (..), Proof (..), Verification (Verified), genItems, prove, readProof, verify, writeProof)
+import Control.Monad (forM, forM_, unless)
+import qualified Data.ByteString as BS
 import Data.Word (Word64)
+import System.Directory (createDirectory, doesDirectoryExist)
 import System.Environment (getArgs)
 import System.Exit (ExitCode (..), exitSuccess, exitWith)
+import System.FilePath ((</>))
 import Test.QuickCheck (generate, resize)
 
 data Command
   = Prove !Options
   | Verify !Options
+  | Generate !Options
 
 data Options = Options
   { size :: Word64
-  , bound :: Word64
   , len :: Int
   , params :: Params
+  , input :: Maybe FilePath
   , output :: FilePath
   }
   deriving (Show)
 
 defaultOptions =
   Options
     { size = 100
-    , bound = 100
     , len = 8
     , params = Params 128 128 80 20
+    , input = Nothing
     , output = "proof.alba"
     }
 
 main :: IO ()
 main = do
   opts <- getArgs >>= parseCommand
   case opts of
-    Prove opts@Options{size, len, params = pars@Params{n_p, n_f}, output} -> do
-      bs <- generate $ resize (fromIntegral size) $ genItems len
-      let opts'@Options{params} = adjustForSize opts
-      putStrLn $ "Generating proof " <> show opts'
-      let prf = prove params bs
-      writeProof output prf >>= \n ->
-        putStrLn ("Written proof to '" <> output <> "' (" <> show n <> " bytes)")
-    Verify opts@Options{size, len, params = pars@Params{n_p, n_f}, output} -> do
-      let opts'@Options{params} = adjustForSize opts
-      putStrLn $ "Verifying proof with " <> show opts'
+    Prove opts@Options{size, len, params, output, input} -> do
+      bs <- case input of
+        Just dir -> do
+          putStrLn $ "Reading items from " <> dir
+          forM [1 .. size] $ \idx ->
+            Bytes <$> BS.readFile (dir </> show idx)
+        Nothing -> do
+          putStrLn "Generating random items"
+          generate $ resize (fromIntegral size) $ genItems len
+      putStrLn $ "Generating proof " <> show opts
+      case prove params bs of
+        Left (NoProof retries) ->
+          putStrLn ("No proof could be written after " <> show retries <> " retries") >> exitWith (ExitFailure 1)
+        Right prf@Proof{retryCount} ->
+          writeProof output prf >>= \n ->
+            putStrLn ("Written proof to '" <> output <> "' (" <> show n <> " bytes, " <> show retryCount <> " retries)")
+    Verify opts@Options{size, len, params, output} -> do
+      putStrLn $ "Verifying proof with " <> show opts
       readProof output >>= \prf ->
         case verify params prf of
           Verified{} -> putStrLn ("Verified proof " <> show prf)
           other -> putStrLn ("Cannot verify proof " <> show prf <> ", failure: " <> show other) >> exitWith (ExitFailure 1)
+    Generate opts@Options{output} -> do
+      let baseDir = output
+      putStrLn $ "Generating random items to " <> baseDir
+      exist <- doesDirectoryExist baseDir
+      unless exist $ createDirectory baseDir
+      bs <- generate $ resize (fromIntegral $ size opts) $ genItems (len opts)
+      forM_ (zip bs [1 ..]) $
+        \(Bytes bytes, idx) -> BS.writeFile (baseDir </> show idx) bytes
+      putStrLn $ "Generated " <> show (length bs) <> " items"
 
 usage :: IO ()
 usage =
@@ -56,31 +78,31 @@ usage =
       [ "alba: Command-line utility for creating and verifying ALBA proofs"
       , ""
       , "Usage:"
-      , "alba prove <options>  : Generate an ALBA proof file from a (random) set of items"
-      , "alba verify <options> : Verify an ALBA proof. Note that options must be consistent with"
-      , "                        the options used for proving"
+      , "alba prove <options>    : Generate an ALBA proof file from a (random) set of items"
+      , "alba verify <options>   : Verify an ALBA proof. Note that options must be consistent with"
+      , "alba generate <options> : Generate a set of random items to be used for proof generation"
       , ""
       , "Options:"
       , "--help           : Display this help text"
       , "--security <int> : The security level of the proof (default: 128)"
-      , "--size <int>     : The actual number of elements to build a proof for (default: bound * honest_ratio)"
-      , "--bound <int>    : The maximum number of elements in the input set (default: 100)"
+      , "--size <int>     : The actual number of elements to build a proof for (default: 100)"
+      , "--n-p <int>      : ALBA n_p parameter, e.g expected \"honest\" set size (default: 80)"
+      , "--n-f <int>      : ALBA n_f parameter, e.g expected \"faulty\" set size (default: 20)"
       , "--len <int>      : The length (in bytes) of each item in the input set (default: 8)"
-      , "--honest-ratio <int>"
-      , "                 : The assumed _percentage_ of \"honest\" items in the input set (default: 80)"
-      , "--output <file>  : The file containing proof to write or verify (default: alba.proof)"
+      , "--output <file>  : The file containing proof to write or verify, or the directory where to generate"
+      , "                   items in (default: alba.proof)"
+      , "--input <dir>    : If set, reads the item to prove from the given directory instead of generating"
+      , "                   them"
       ]
 
-adjustForSize :: Options -> Options
-adjustForSize opts@Options{size, bound, params = pars@Params{n_p, n_f}} =
-  opts{params = pars{n_p = bound * n_p `div` 100, n_f = bound * n_f `div` 100}}
-
 parseCommand :: [String] -> IO Command
 parseCommand = \case
   ("prove" : rest) ->
     Prove <$> parseOptions rest
   ("verify" : rest) ->
     Verify <$> parseOptions rest
+  ("generate" : rest) ->
+    Generate <$> parseOptions rest
   ("--help" : _) ->
     usage >> exitSuccess
   other ->
@@ -89,7 +111,7 @@ parseCommand = \case
 parseOptions = \case
   [] -> pure defaultOptions
   ("--help" : _) ->
-    usage >> exitWith ExitSuccess
+    usage >> exitSuccess
   ("--security" : lam : rest) -> do
     let λ = read lam
     opts <- parseOptions rest
@@ -98,20 +120,23 @@ parseOptions = \case
     let size = read sz
     opts <- parseOptions rest
     pure $ opts{size}
-  ("--bound" : sz : rest) -> do
-    let bound = read sz
+  ("--n-p" : np : rest) -> do
+    let n_p = read np
+    opts <- parseOptions rest
+    pure $ opts{params = (params opts){n_p}}
+  ("--n-f" : nf : rest) -> do
+    let n_f = read nf
     opts <- parseOptions rest
-    pure $ opts{bound}
+    pure $ opts{params = (params opts){n_f}}
   ("--len" : ln : rest) -> do
     let len = read ln
     opts <- parseOptions rest
     pure $ opts{len}
-  ("--honest-ratio" : hn : rest) -> do
-    let rat = read hn
-    opts <- parseOptions rest
-    pure $ opts{params = (params opts){n_p = rat, n_f = 100 - rat}}
   ("--output" : output : rest) -> do
     opts <- parseOptions rest
     pure $ opts{output}
+  ("--input" : input : rest) -> do
+    opts <- parseOptions rest
+    pure $ opts{input = Just input}
   other -> do
     usage >> exitWith (ExitFailure 2)