content: module Test.QAP where
  
  import Protolude
  
  import           Data.Map              (Map)
  import qualified Data.Map              as Map
  import           Test.Tasty.HUnit
  import           Test.Tasty.QuickCheck
  
  import Circuit.Affine
  import Circuit.Arithmetic
  import Data.Pairing.BN254 (Fr, getRootOfUnity)
  import QAP
  
  import Test.Circuit.Affine (arbAffineCircuit, arbInputVector)
  
  -------------------------------------------------------------------------------
  -- Generators
  -------------------------------------------------------------------------------
  
  arbGate
    :: Arbitrary f
    => Int -> Gen (Gate Wire f)
  arbGate numVars
    = oneof
      [ Mul <$> (mapVarsAffine InputWire <$> sized (arbAffineCircuit numVars))
            <*> (mapVarsAffine InputWire <$> sized (arbAffineCircuit numVars))
            <*> pure (OutputWire 0)
      , Equal <$> (InputWire <$> choose (0, numVars - 1))
              <*> pure (IntermediateWire 0)
              <*> pure (OutputWire 0)
      ]
  
  data GateWithInputs f = GateWithInputs (Gate Wire f) [Map Int f]
    deriving Show
  
  instance Arbitrary f => Arbitrary (GateWithInputs f) where
    arbitrary = do
      numVars <- (+1) . abs <$> arbitrary
      program <- scale (`div` 7) $ arbGate numVars
      inputs <- vectorOf 10 (arbInputVector numVars)
      pure $ GateWithInputs program inputs
  
  -------------------------------------------------------------------------------
  -- Test values
  -------------------------------------------------------------------------------
  
  testArithCircuit :: ArithCircuit Fr
  testArithCircuit
    = ArithCircuit
      [ Mul (Var (InputWire 0)) (Var (InputWire 1)) (IntermediateWire 0)
      , Mul (Var (InputWire 2)) (Var (InputWire 3)) (IntermediateWire 1)
      , Mul (Add (ConstGate 10) (Var (IntermediateWire 0))) (Var (IntermediateWire 1)) (OutputWire 0)
      ]
  
  testVarsArithCircuit :: Map Int Fr
  testVarsArithCircuit
    = Map.fromList [ (0, 2)
                   , (1, 3)
                   , (2, 4)
                   , (3, 5)
                   ]
  
  -------------------------------------------------------------------------------
  -- Tests
  -------------------------------------------------------------------------------
  
  unit_arithCircuitToQapCorrect :: Assertion
  unit_arithCircuitToQapCorrect
    = assertBool "Verifying assignment against QAP of circuit failed"
    $ verifyAssignment qap assignment
    where
      roots = [[7],[8],[9]]
      qap = arithCircuitToQAP roots testArithCircuit
      assignment = generateAssignment testArithCircuit testVarsArithCircuit
  
  unit_arithCircuitToQapNoFalsePositive :: Assertion
  unit_arithCircuitToQapNoFalsePositive
      = assertBool "Verifying whether verification fails on faulty assignment"
      $ not $ verifyAssignment qap invalidAssignment
    where
      roots = [[7],[8],[9]]
      qap = arithCircuitToQAP roots testArithCircuit
      invalidAssignment
        = QapSet
          { qapSetConstant = 1
          , qapSetInput = Map.fromList [(0,2),(1,3),(2,4),(3,5)]
          , qapSetIntermediate = Map.fromList [(0,7),(1,20)]
          , qapSetOutput = Map.fromList [(0,320)]
          }
  
  prop_gateToQapCorrect
    :: GateWithInputs Fr -> Bool
  prop_gateToQapCorrect (GateWithInputs program inputs)
    = all testInput inputs
    where
      roots = case program of
               Mul {} -> [1]
               Equal {} -> [1,2]
               _ -> panic "Invalid roots"
      qap = gateToQAP getRootOfUnity roots program
      assignment input = generateAssignmentGate program input
      testInput = verifyAssignment qap . assignment
  