cabal-version:      2.4
name:               crypto-junk
version:            0.1.0.0
synopsis:           Half-completed crypto experiments in Haskell.
-- description:
homepage:           https://nest.pijul.com/sellout/crypto-junk
bug-reports:        https://nest.pijul.com/sellout/crypto-junk/discussions
license:            AGPL-3.0-or-later
author:             Greg Pfeil
maintainer:         greg@technomadic.org
copyright:          2022 Greg Pfeil
category:           Cryptography
extra-source-files: CHANGELOG.md
library
    exposed-modules:  Halo2.PLONKish
    ghc-options:     -Weverything
    build-depends:    base
                    , containers
                    , either
                    , fin
                    , vec
    default-language: Haskell2010

{-# language
        ConstraintKinds,
        DataKinds,
        DeriveDataTypeable,
        DeriveGeneric,
        DeriveTraversable,
        GADTs,
        ImportQualifiedPost,
        LambdaCase,
        RankNTypes,
        StandaloneDeriving
#-}
module Halo2.PLONKish
  ( OrderedPair,
    orderedPair,
    MultivarPoly (..),
    MultivariatePolynomial (..),
    evaluateMultivariatePolynomial,
    ColumnSpecification (..),
    Configuration (..),
    Circuit (configuration, values),
    circuit,
    EqualityConstraintFailure (..),
    addEqualityConstraint,
    -- generateProvingKey,
    -- generateVerificationKey,
  )
where
import Control.Exception (Exception (..))
import Data.Data (Data)
import Data.Either.Combinators (maybeToLeft)
import Data.Fin (Fin)
import Data.Foldable (toList)
import Data.List.NonEmpty (NonEmpty, nonEmpty)
import Data.Map (Map)
import Data.Map qualified as Map
import Data.Maybe (catMaybes)
import Data.Proxy (Proxy (..))
import Data.Semigroup (Sum (..))
import Data.Set (Set)
import Data.Set qualified as Set
import Data.Type.Nat (Nat (..))
import Data.Type.Nat qualified as Nat
import Data.Typeable (Typeable)
import Data.Vec.Lazy (Vec)
import Data.Vec.Lazy qualified as Vec
import GHC.Generics (Generic, Generic1)
-- | A pair that maintains keeps the elements in increasing order
data OrderedPair a = OrderedPair a a
  deriving (Data, Eq, Generic, Ord, Foldable, Functor, Generic1, Traversable)
orderedPair :: Ord a => a -> a -> OrderedPair a
orderedPair a b = if a < b then OrderedPair a b else OrderedPair b a
-- | A sequence with /at most/ @n@ elements.
data BoundedList n a where
  BNil :: BoundedList n a
  BCons :: a -> BoundedList n a -> BoundedList ('S n) a
  -- deriving (Data, Generic, Ord, Foldable, Functor, Generic1, Traversable)
deriving instance Eq a => Eq (BoundedList n a)
deriving instance Foldable (BoundedList n)
-- | Make a smaller `BoundedList` fit into a space that requires a larger one.
weaken ::
  forall m n a. Proxy m -> BoundedList n a -> BoundedList (Nat.Plus n m) a
weaken = go where
    go :: Proxy m -> BoundedList o a -> BoundedList (Nat.Plus o m) a
    go Proxy BNil = BNil
    go m (BCons x xs) = BCons x (go m xs)
-- | A `BoundedList` whose elements are maintained in order.
newtype OrderedBoundedList n a =
  OrderedBoundedList { unorderedBoundedList :: BoundedList n a }
  deriving (Foldable)
bcons :: a -> OrderedBoundedList n a -> OrderedBoundedList ('S n) a
bcons a = OrderedBoundedList . BCons a . unorderedBoundedList
type Field = Num
-- | Represents a multivariate polynomial with a maximum degree.
--
--   This breaks down the exponents into repetitions of the variables. E.g..,
--   @7 * x^2 * y * z^3@ has the variables mapped to x=0, y=1, z=2 and produces
--   a `Map` entry that looks like @([0, 0, 1, 2, 2, 2], 7)`. The list is
--   maintained in order and guaranteed (by the type system) to never have more
--   than @maximumDegree@ elements.
type MVP vars maximumDegree field =
  Map (OrderedBoundedList maximumDegree (Fin vars)) field
-- | Lowers a multivariate polynomial to a function.
evaluateMVP ::
  Field field => MVP vars maximumDegree field -> Vec vars field -> field
evaluateMVP mvp values =
  getSum $
    Map.foldMapWithKey
      ( \k coefficient ->
          Sum $ foldr (\i acc -> (values Vec.! i) * acc) coefficient k
      )
      mvp
data MultivariatePolynomial f maximumDegree c vars
     -- |
     --
     --  __NB__: We could keep @vars@ universal here, but we'd then have to wrap
     --          it in another @data@ declaration in order to use it in other
     --          structures. We shuld re-evaluate this later.
  = MultivariatePolynomial
  { -- | The row here is `Integer`, because in the `Configuration` we don't know
    --   how many rows we have, and so the actual row is calculated later,
    --   relative to the current row, modulo the number of rows.
    variables :: Vec vars (Fin c, Integer),
    polynomial :: MVP vars maximumDegree f
  }
evaluateMultivariatePolynomial
  :: (Field f, Nat.SNatI r)
  => MultivariatePolynomial f maximumDegree c vars
  -> Circuit f c r maximumDegree
  -> Fin r
  -> f
evaluateMultivariatePolynomial mvp circuit currentRow =
  evaluateMVP (polynomial mvp)
    . fmap
      ( \(c, r) ->
          (values circuit Vec.! c) Vec.! (currentRow + fromInteger r)
      )
    $ variables mvp
data ColumnSpecification
  = Fixed
  | Advice
  | Instance
data MultivarPoly f maximumConstraintDegree c =
  forall v. MultivarPoly (MultivariatePolynomial f maximumConstraintDegree c v)
-- @f@ should be a `FiniteField`, @c@ is the number of columns.
data Configuration f c maximumConstraintDegree = Configuration
  { columns :: Vec c ColumnSpecification,
    equalityConstraintColumns :: Set (Fin c),
    polynomialConstraints :: Set (MultivarPoly f maximumConstraintDegree c),
    -- |
    --
    --  __TODO__: Determine if a `Map` is the right structure here. I.e., do we
    --            ever need to look up by column and do we need to support
    --            duplicate polynomials? (If the latter, we need to decide
    --            between a `MultiMap` and a @`Map` (`MultivariatePolynomial` f)
    --            [`Fin` c]@).
    lookupArguments :: Map (MultivarPoly f maximumConstraintDegree c) (Fin c)
  }
checkEqualityColumn :: Configuration f c mcd -> Fin c -> Bool
checkEqualityColumn = flip Set.member . equalityConstraintColumns
-- | @r@ is the number of rows in the `Circuit`.
data Circuit f c r maximumConstraintDegree = Circuit
  { configuration :: Configuration f c maximumConstraintDegree,
    -- |
    --
    --  __NB__: `Set` maintains its elements in ascending order, so we do have
    --          determinism for generating keys.
    equalityConstraints' :: Set (OrderedPair (Fin c, Fin r)),
    values :: Vec c (Vec r f)
  }
circuit
  :: Configuration f c maximumConstraintDegree
  -> Vec c (Vec r f)
  -> Circuit f c r maximumConstraintDegree
circuit conf = Circuit conf Set.empty
data EqualityConstraintFailure c r
  = IllegalEqualityColumn (Fin c) (Set (Fin c))
    -- |
    --
    --  __TODO__: This should actually be a warning, not an error.
  | TautologicalEqualityConstraint (Fin c, Fin r)
  deriving (Eq, Ord)
-- |
--
--  __TODO__: This should not be defined here, and the output should be
--            rephrased to make sense in the context in which it's forced.
instance Show (EqualityConstraintFailure c r) where
  show = \case
    IllegalEqualityColumn column allowedColumns ->
      "Column "
        <> show column
        <> " is not in the set of columns that can participate in equality constraints (`Halo2.PLONKish.Configuration.equalityConstraintColumns`): "
        <> show allowedColumns
        <> "."
    -- Documenting this as if it were a warning, even though it currently causes
    -- a failure.
    TautologicalEqualityConstraint cell ->
      "The equality constraint "
        <> show cell
        <> " == "
        <> show cell
        <> "is tautological, because both elements are the same. Refusing to add it."
instance (Typeable c, Typeable r) => Exception (EqualityConstraintFailure c r)
addEqualityConstraint
  :: (Fin c, Fin r)
  -> (Fin c, Fin r)
  -> Circuit f c r mcd
  -> Either (NonEmpty (EqualityConstraintFailure c r)) (Circuit f c r mcd)
addEqualityConstraint a b circuit =
  maybeToLeft
    ( circuit {
        equalityConstraints' =
          Set.insert (orderedPair a b) (equalityConstraints' circuit)
      }
    )
    . nonEmpty
    $ validateEqualityConstraint (configuration circuit) a b 
validateEqualityConstraint
  :: Configuration f c mcd
  -> (Fin c, Fin r)
  -> (Fin c, Fin r)
  -> [EqualityConstraintFailure c r]
validateEqualityConstraint config a b =
  let invalidColumns =
        toList . Set.fromList $
           catMaybes [reifyEqualityCheck config a, reifyEqualityCheck config b]
   in if a == b
        then TautologicalEqualityConstraint a : invalidColumns
        else invalidColumns
  where
    reifyEqualityCheck config x =
      if checkEqualityColumn config (fst x)
        then Nothing
        else
          Just . IllegalEqualityColumn (fst x) $
            equalityConstraintColumns config
-- | This builds an elliptic curve from the two coefficients, returning a
--   function that, given the x component of a point will return the positive y
--   component. The negative component can be found via point inversion.
ellipticCurve :: Field f => f -> f -> f -> f
ellipticCurve a b x = sqrt $ x ** 3 + a * x + b
type F_101 = ()
meh :: F_101 -> F_101
meh = ellipticCurve 0 3
type ProvingKey = ()
type VerificationKey = ()
generateProvingKey :: Circuit f c r mcd -> ProvingKey
generateProvingKey = _
generateVerificationKey :: Circuit f c r mcd -> VerificationKey
generateVerificationKey = _