content: module Test.Circuit.Expr where
  
  import           Circuit.Arithmetic
  import           Circuit.Expr
  import           Data.Curve.Weierstrass.BN254 (Fr)
  import qualified Data.Map                     as Map
  import           Fresh
  import           Protolude
  import           QAP
  import           Test.Circuit.Affine
  import           Test.Tasty.QuickCheck
  
  -------------------------------------------------------------------------------
  -- Generators
  -------------------------------------------------------------------------------
  
  arbExprBool :: Arbitrary f => Int -> Int -> Gen (Expr Int f Bool)
  arbExprBool numVars size
    | size <= 0 = oneof $ [EConstBool <$> arbitrary] ++ if numVars > 0
      then []
      else []
    | size > 0 = oneof
      [ EBinOp BAnd <$> arbExprBool numVars (size - 1) <*> arbExprBool
        numVars
        (size - 1)
      , EBinOp BOr <$> arbExprBool numVars (size - 1) <*> arbExprBool numVars
                                                                      (size - 1)
      , EUnOp UNot <$> arbExprBool numVars (size - 1)
      , EEq <$> arbExpr numVars (size - 1)
            <*> arbExpr numVars (size - 1)
      ]
  
  arbExpr :: Arbitrary f => Int -> Int -> Gen (Expr Int f f)
  arbExpr numVars size
    | size <= 0 = oneof $ [EConst <$> arbitrary] ++ if numVars > 0
      then [EVar <$> choose (0, numVars - 1)]
      else []
    | size > 0 = oneof
      [ EBinOp BAdd <$> arbExpr numVars (size - 1) <*> arbExpr numVars (size - 1)
      , EBinOp BSub <$> arbExpr numVars (size - 1) <*> arbExpr numVars (size - 1)
      , EBinOp BMul <$> arbExpr numVars (size - 1) <*> arbExpr numVars (size - 1)
      , EUnOp UNeg <$> arbExpr numVars (size - 1)
      , EIf
      <$> arbExprBool numVars (size - 1)
      <*> arbExpr     numVars (size - 1)
      <*> arbExpr     numVars (size - 1)
      ]
  
  
  data ExprWithInputs f = ExprWithInputs (Expr Int f f) [Map Int f]
    deriving Show
  
  
  instance Arbitrary f => Arbitrary (ExprWithInputs f) where
    arbitrary = do
      numVars <- abs <$> arbitrary
      program <- scale (`div` 10) $ sized (arbExpr numVars)
      inputs <- vectorOf 5 $ arbInputVector numVars
      pure $ ExprWithInputs program inputs
  
  -------------------------------------------------------------------------------
  -- Tests
  -------------------------------------------------------------------------------
  
  
  -- | Check whether exprToArithCircuit produces valid circuits
  prop_compiledCircuitValid :: ExprWithInputs Fr -> Bool
  prop_compiledCircuitValid (ExprWithInputs expr _) =
    validArithCircuit (execCircuitBuilder $ exprToArithCircuit expr (OutputWire 0))
  
  -- | Check whether exprToArithCircuit produces circuits that have valid assignments to the qap
  prop_compiledQAPValid :: ExprWithInputs Fr -> Property
  prop_compiledQAPValid (ExprWithInputs expr inputs) = withMaxSuccess 50
    $ all testInput inputs
   where
    circuit = (execCircuitBuilder $ exprToArithCircuit expr (OutputWire 0))
    roots :: ArithCircuit Fr -> [[Fr]]
    roots = evalFresh . generateRoots (fromIntegral <$> fresh)
    qap = arithCircuitToQAP (roots circuit) circuit
    assignment input = generateAssignment circuit input
    testInput = verifyAssignment qap . assignment
  
  -- | Check whether evaluating an expression and
  -- evaluating the arithmetic circuit translation produces the same
  -- result
  prop_evalEqArithEval :: ExprWithInputs Fr -> Bool
  prop_evalEqArithEval (ExprWithInputs expr inputs) = all testInput inputs
   where
    testInput input = exprResult input == arithResult input
    exprResult input = evalExpr (Map.lookup) expr input
    arithResult input = arithOutput input Map.! (OutputWire 0)
    arithOutput input = evalArithCircuit (Map.lookup)
                                         (Map.insert)
                                         circuit
                                         (Map.mapKeys InputWire input)
    circuit = (execCircuitBuilder $ exprToArithCircuit expr (OutputWire 0))
  