content: module Test.Circuit.Affine where
  
  import           Circuit.Affine
  import qualified Data.Map              as Map
  import           Protolude
  import           Test.Tasty.QuickCheck
  
  -------------------------------------------------------------------------------
  -- Generators
  -------------------------------------------------------------------------------
  
  arbAffineCircuit ::
    Arbitrary f =>
    Int ->
    Int ->
    Gen (AffineCircuit Int f)
  arbAffineCircuit numVars size
    | size <= 0 =
      oneof $
        [ ConstGate <$> arbitrary
        ]
          ++ if numVars > 0
            then [Var <$> choose (0, numVars - 1)]
            else []
    | size > 0 =
      oneof
        [ ScalarMul <$> arbitrary <*> arbAffineCircuit numVars (size - 1),
          Add <$> arbAffineCircuit numVars (size - 1)
            <*> arbAffineCircuit numVars (size - 1)
        ]
  
  arbInputVector :: Arbitrary f => Int -> Gen (Map Int f)
  arbInputVector numVars = Map.fromList . zip [0 ..] <$> vector numVars
  
  -- | The input vector has to have the correct length, so we want to
  -- generate the program and the test input simultaneously.
  data AffineCircuitWithInputs f = AffineCircuitWithInputs (AffineCircuit Int f) [Map Int f]
    deriving (Show)
  
  instance Arbitrary f => Arbitrary (AffineCircuitWithInputs f) where
    arbitrary = do
      numVars <- abs <$> arbitrary
      program <- scale (`div` 7) $ sized (arbAffineCircuit numVars)
      inputs <- vectorOf 10 $ arbInputVector numVars
      pure $ AffineCircuitWithInputs program inputs
  
  -------------------------------------------------------------------------------
  -- Tests
  -------------------------------------------------------------------------------
  
  -- | Check that evaluating the vector representation of the circuit
  -- yields the same results as evaluating the circuit "directly". Field
  -- is instantiated as being the rationals for testing. It later should
  -- probably be something like Pairing.Fr.Fr.
  prop_affineCircuitToAffineMap ::
    AffineCircuitWithInputs Rational -> Bool
  prop_affineCircuitToAffineMap (AffineCircuitWithInputs program inputs) =
    all testInput inputs
    where
      testInput input =
        evalAffineCircuit Map.lookup input program
          == evalAffineMap (affineCircuitToAffineMap program) input
  