BLS12381.hs

-- SPDX-FileCopyrightText: 2021 Oxhead Alpha
-- SPDX-License-Identifier: LicenseRef-MIT-OA

-- | Support for BLS12-381 elliptic curve.
--
-- Some general hints on the implementation can be found in this python
-- re-implementation used by Tezos for testing:
-- <https://gitlab.com/metastatedev/tezos/-/commit/f10c39e0030e6b4fdd416a62de7b80b6ffdfeacf#80b4b1585c1e6fa82f2cfaf75001c490613f22c3>.
-- And it uses this library inside: <https://github.com/ethereum/py_ecc/tree/master/py_ecc/optimized_bls12_381>.
module Morley.Tezos.Crypto.BLS12381
  ( Bls12381Fr
  , Bls12381G1
  , Bls12381G2
  , CurveObject (..)
  , MultiplyPoint (..)
  , DeserializationError (..)
  , checkPairing

    -- * Playground
  , unsafeReadFromHex
  , generateFrom
  , g1One
  , g2One
  ) where

import Prelude hiding (negate, one)
import Prelude qualified
import Unsafe qualified (fromIntegral)

import Control.Exception (assert)
import Control.Lens (each, toListOf)
import Control.Monad.Random (MonadRandom, evalRand, getRandom, mkStdGen)
import Data.Bits (bit, complement, setBit, testBit, (.&.))
import Data.ByteString qualified as BS
import Data.Curve qualified as C
import Data.Curve.Weierstrass qualified as CW
import Data.Curve.Weierstrass.BLS12381 qualified as CW.BLS
import Data.Field.Galois qualified as GF
import Data.Pairing.BLS12381 qualified as BLS
import Fmt (Buildable(..), hexF, (<+>))
import Morley.Util.Instances ()
import Morley.Util.Named
import Text.Hex (decodeHex)


-- | Methods common for all BLS12-381 primitives.
class CurveObject a where
  -- | Representation of @0@, aka additive identity.
  zero :: a

  -- | Negate a value.
  negate :: a -> a

  -- | Add up two values.
  add :: a -> a -> a

  -- | Generate a random value.
  generate :: MonadRandom m => m a

  -- | Read a value from Michelson's bytes form.
  --
  -- Michelson tends to represent all BLS12-381 types in bytes form,
  -- some special types also allow other forms.
  fromMichelsonBytes :: ByteString -> Either DeserializationError a

  -- | Produce Michelson's bytes representation.
  toMichelsonBytes :: a -> ByteString

-- | Generate a random value from given seed.
generateFrom :: (CurveObject a) => Int -> a
generateFrom = evalRand generate . mkStdGen

-- | Reads an object from hex string.
--
-- To be used only in playground and tests.
unsafeReadFromHex :: (CurveObject a, HasCallStack) => String -> a
unsafeReadFromHex hex =
  let bs = decodeHex (toText hex) ?: error "bad hex"
  in unsafe $ fromMichelsonBytes bs

-- | Multiplication operations on BLS12-381 objects.
class MultiplyPoint scalar point where
  -- | Multiply point value by scalar value.
  multiply :: scalar -> point -> point

-- | G1 point on the curve.
newtype Bls12381G1 = Bls12381G1 { unBls12381G1 :: BLS.G1' }
  deriving stock (Show, Eq)
  deriving newtype (NFData)

instance CurveObject Bls12381G1 where
  zero = Bls12381G1 C.id
  negate (Bls12381G1 v) = Bls12381G1 (C.inv v)
  add (Bls12381G1 a) (Bls12381G1 b) = Bls12381G1 (C.add a b)
  generate = Bls12381G1 <$> C.rnd
  fromMichelsonBytes =
    let bsToCoord = toPrime . fromBigEndian
    in fmap Bls12381G1 . parseJA2WAPoint g1CoordLen bsToCoord
  toMichelsonBytes =
    let coordToBs = toBigEndian g1CoordLen . fromPrime
    in convertWA2JAPoint g1CoordLen coordToBs . unBls12381G1

instance MultiplyPoint Integer Bls12381G1 where
  multiply s (Bls12381G1 p) = Bls12381G1 (C.mul' p s)

-- | G2 point on the curve.
newtype Bls12381G2 = Bls12381G2 { unBls12381G2 :: BLS.G2' }
  deriving stock (Show, Eq)
  deriving newtype (NFData)

instance CurveObject Bls12381G2 where
  zero = Bls12381G2 C.id
  negate (Bls12381G2 v) = Bls12381G2 (C.inv v)
  add (Bls12381G2 a) (Bls12381G2 b) = Bls12381G2 (C.add a b)
  generate = Bls12381G2 <$> C.rnd
  fromMichelsonBytes =
    let fromBsPair = map fromBigEndian . toListOf each . BS.splitAt (g2CoordLen `div` 2)
        bsToCoord = GF.toE . reverse . map toPrime . fromBsPair
    in fmap Bls12381G2 . parseJA2WAPoint g2CoordLen bsToCoord
  toMichelsonBytes =
    let toBsPair = foldMap (toBigEndian $ g2CoordLen `div` 2)
        coordToBs = toBsPair . map fromPrime . reverse . GF.fromE
    in convertWA2JAPoint g1CoordLen coordToBs . unBls12381G2

instance MultiplyPoint Integer Bls12381G2 where
  multiply s (Bls12381G2 p) = Bls12381G2 (C.mul' p s)

-- | An element of an algebraic number field (scalar), used for multiplying
-- 'Bls12381G1' and 'Bls12381G2'.
newtype Bls12381Fr = Bls12381Fr { unBls12381Fr :: BLS.Fr }
  deriving stock (Show, Eq, Ord)
  deriving newtype (Num, Enum, Bounded, Real, Integral, Fractional, NFData)

instance CurveObject Bls12381Fr where
  zero = Bls12381Fr 0
  negate = Prelude.negate
  add = (+)
  generate = Bls12381Fr <$> getRandom
  fromMichelsonBytes bs =
    if length bs > frLen
    then Left $ TooLargeLength ! #limit frLen ! #given (length bs)
    else
      let num = fromLittleEndian bs
      in fromIntegralNoOverflow num
           & first (\_ -> ValueOutsideOfField $ toInteger num)
  toMichelsonBytes = toLittleEndian frLen . fromPrime . unBls12381Fr

instance MultiplyPoint Bls12381Fr Bls12381G1 where
  multiply (Bls12381Fr s) (Bls12381G1 p) = Bls12381G1 (C.mul p s)

instance MultiplyPoint Bls12381Fr Bls12381G2 where
  multiply (Bls12381Fr s) (Bls12381G2 p) = Bls12381G2 (C.mul p s)

-- | Checks that product of pairings of points in the list is equal to 1 in
-- Fq12 field.
checkPairing :: [(Bls12381G1, Bls12381G2)] -> Bool
checkPairing pairs =
  -- Some hints for implementation of this function:
  -- https://gitlab.com/metastatedev/tezos/-/commit/bb2cda17d48a52ce854e027f0222a0463e0e66f0#af97cb649204420968454a94e7bfaa6a6e27195a_1285_1330
  -- https://gitlab.com/metastatedev/tezos/-/commit/f10c39e0030e6b4fdd416a62de7b80b6ffdfeacf#80b4b1585c1e6fa82f2cfaf75001c490613f22c3_0_172

  -- Monoid instance on GT' has the desired multiplicative semantics
  foldMap pairing pairs == mempty
  where
    pairing :: (Bls12381G1, Bls12381G2) -> BLS.GT'
    pairing (Bls12381G1 g1, Bls12381G2 g2) =
      BLS.finalExponentiationBLS12 BLS.parameterHex
        (BLS.millerAlgorithmBLS12 BLS.parameterBin g1 g2)

----------------------------------------------------------------------------
-- Serialization helpers
----------------------------------------------------------------------------

-- | All kinds of errors that can occur when reading a Michelson value.
data DeserializationError
  = CompressedFormIsNotSupported
  | UnexpectedLength ("expected" :! Int) ("given" :! Int)
  | TooLargeLength ("limit" :! Int) ("given" :! Int)
  | ValueOutsideOfField Integer
  | PointNotOnCurve ByteString
  deriving stock (Show, Eq, Generic)
  deriving anyclass (NFData)

instance Buildable DeserializationError where
  build = \case
    CompressedFormIsNotSupported ->
      "Compressed form of BLS12-381 point is not supported by Tezos"
    UnexpectedLength (arg #expected -> expected) (arg #given -> given) ->
      "Unexpected length of BLS12-381 primitive: \
      \expected" <+> build expected <> ", but given" <+> build given
    TooLargeLength (arg #limit -> limit) (arg #given -> given) ->
      "Too large length of BLS12-381 primitive: \
      \limit is" <+> build limit <+> ", but given" <+> build given
    ValueOutsideOfField v ->
      "Value is too large for the given field of values:" <+> build v
    PointNotOnCurve bs ->
      "Point is not on curve: 0x" <> hexF bs

{- Note on serialization:

In BLS12-381, "381" stands for the number of bits necessary to represent a
coordinate of a point on a curve, i.e. we have to use 48 bytes, getting 3 extra
bits. In the BLS12-381 library used by Tezos, those bits are exploited to carry
some meta information.

More info can be found here:
<https://github.com/zkcrypto/pairing/blob/307aa1f29dccaed09abe774d2027cad57fc5d0b4/src/bls12_381/README.md#serialization>.

So Fr is just a scalar (but a pretty big one), represented in little-endian
as said in the Michelson docs.

G1 and G2 represent a point on curve and have the following form:

      X coordinate         Y coordinate
 |!___________________|____________________|
   \  (big-endian)         (big-endian)
    \
     `- bits with meta info

Generally, various coordinate systems may be used to represent a point on curve,
and the library picked by Tezos uses Jacobian coordinates, probably assuming
that the third @Z@ coordinate is always @1@.
Note that we use 'Data.Pairing.BLS12381' which by default picks Weierstrass
coordinates (it has two, not three coordinates, and an "infinity point" which
is kept as a special case), but it also provides methods for converting between
different coordinate systems.

Coordinates take a different amount of space in G1 and G2:
* In G1 both X and Y are from a so-called "Fr" field, where numbers take
  48 bytes (without the leading 3 bits).
* In G2 both coordinates are from "Fr2" field which is a two-dimensional field
  over "Fr", i.e. X and Y themselves contain two 48-byte coordinates each.

Note that this is correct for the "uncompressed" form, and there is a different
"compressed" form that, fortunately, seems to be not supported by Michelson.
They initially planned to have @COMPRESS@ and @UNCOMPRESS@ instructions, perhaps
for manual conversions, but those instructions didn't appear in Edonet
eventually.

-}

-- | A helper datatype for representing points in raw bytes form.
data RawPoint
  = Infinity
    -- ^ Point at infinity.
  | RawPoint ByteString
    -- ^ Bytes representing the payload.

-- | Given the Michelson representation of a point, interpret flags
-- and return the bare payload.
--
-- This assumes that a proper number of bytes is provided.
parsePointFlags
  :: HasCallStack
  => ByteString -> Either DeserializationError RawPoint
parsePointFlags bsFull =
  case BS.uncons bsFull of
    Nothing -> error "Empty byte sequence"
    Just (b, bs)
      | b `testBit` compressionBit ->
          Left CompressedFormIsNotSupported
      | b `testBit` infinityBit ->
          return Infinity
      | otherwise -> do
          let
            b' = b .&. complement
              (sum $ map bit [compressionBit, infinityBit, flag3Bit])
          return $ RawPoint (BS.cons b' bs)

-- | Fill a point in raw bytes form with the necessary flags.
fillPointFlags :: HasCallStack => Int -> RawPoint -> ByteString
fillPointFlags 0 = error "Coordinates are unexpectedly empty"
fillPointFlags len = \case
  Infinity -> BS.cons (0 `setBit` infinityBit) (BS.replicate (len - 1) 0)
  RawPoint bs -> bs

-- | Get a bytestring containing coordinates of a point and split it,
-- checking that each coordinate occupies the given number of bytes.
splitUncompressedPoint
  :: Int -> ByteString -> Either DeserializationError (ByteString, ByteString)
splitUncompressedPoint coordLen bs
  | length bs /= coordLen * 2 =
      Left $ UnexpectedLength ! #expected (coordLen * 2) ! #given (length bs)
  | otherwise =
      Right $ BS.splitAt coordLen bs

-- | Parse a point in Weierstrass form and Affine coordinates,
-- assuming that in the provided bytestring the point is given in Jacobian
-- coordinates (the library used by Tezos operates in Jacobian coordinates).
parseJA2WAPoint
  :: ( CW.BLS.WJCurve CW.BLS.BLS12381 fq BLS.Fr
     , CW.BLS.WACurve CW.BLS.BLS12381 fq BLS.Fr
     )
  => Int
  -> (ByteString -> fq)
  -> ByteString
  -> Either DeserializationError (CW.BLS.WAPoint CW.BLS.BLS12381 fq BLS.Fr)
parseJA2WAPoint coordLen toCoord full = do
  (xRawWithFlags, yRaw) <- splitUncompressedPoint coordLen full
  xRawPoint <- parsePointFlags xRawWithFlags
  case xRawPoint of
    Infinity -> return CW.O
    RawPoint xRaw ->
      let point = C.toA $ CW.J (toCoord xRaw) (toCoord yRaw) 1
      in if C.def point
         then return point
         else Left $ PointNotOnCurve full

-- | Turn a Weierstrass Affine point into Jacobian coordinates and represent
-- those as bytes.
convertWA2JAPoint
  :: ( CW.BLS.WJCurve CW.BLS.BLS12381 fq BLS.Fr
     , CW.BLS.WACurve CW.BLS.BLS12381 fq BLS.Fr
     )
  => Int
  -> (fq -> ByteString)
  -> CW.BLS.WAPoint CW.BLS.BLS12381 fq BLS.Fr
  -> ByteString
convertWA2JAPoint coordLen toRawCoord point =
  let
    rawPoint = case point of
      CW.O -> Infinity
      p@CW.A{} -> RawPoint $
        let CW.J x y z = C.fromA p
           -- Conversion from Affine coordinates used to produce an already
           -- normalized value.
           -- In case this turns out to be incorrect, probably @x / z@ and @y / z@
           -- are just what we want.
        in assert (z == 1) $
            toRawCoord x <> toRawCoord y
  in fillPointFlags (coordLen * 2) rawPoint

-- | Interpret a byte sequence as a number in big-endian.
fromBigEndian :: ByteString -> Natural
fromBigEndian bs =
  foldl' (\acc byte -> acc * 0x100 + fromIntegral @Word8 @Natural byte) 0 $
  BS.unpack bs

-- | Interpret a byte sequence as a number in little-endian.
fromLittleEndian :: ByteString -> Natural
fromLittleEndian bs =
  sum . zipWith (*) (iterate (* 0x100) 1) . map (fromIntegral @Word8 @Natural) $
  BS.unpack bs

-- | Represent a number in a big-endian byte sequence, padding the output
-- to the expected length.
--
-- We assert that the length is sufficient for representing the given number.
toBigEndian :: Int -> Natural -> ByteString
toBigEndian len num =
  BS.pack $
  let
    (remainder, bytes) = mapAccumR
      (\x _ -> second (Unsafe.fromIntegral @Natural @Word8) $ x `divMod` 0x100)
      num [1 .. len]
  in assert (remainder == 0) bytes

-- | Represent a number in a little-endian byte sequence, padding the output
-- to the expected length.
--
-- We assert that the length is sufficient for representing the given number.
toLittleEndian :: Int -> Natural -> ByteString
toLittleEndian len num =
  BS.pack $
  let
    (remainder, bytes) = mapAccumL
      (\x _ -> second (Unsafe.fromIntegral @Natural @Word8) $ x `divMod` 0x100)
      num [1 .. len]
  in assert (remainder == 0) bytes

-- | Turn a prime field element into a natural.
fromPrime :: forall p. KnownNat p => GF.Prime p -> Natural
fromPrime p =
  -- This should be safe, since 'Prime's exist in modular arithmetics,
  -- so its conversion to an integer should produce non-negative elements.
  -- In fact, 'GF.Prime' is a newtype wrapper over 'Natural', but
  -- its constructor is not exported :/
  assert (p >= 0) $
    Unsafe.fromIntegral @(GF.Prime p) @Natural p

-- | The inverse to 'fromPrime'.
toPrime :: KnownNat p => Natural -> GF.Prime p
toPrime = fromIntegralOverflowing

-- Primitives' lengths
----------------------------------------------------------------------------

-- | Length of a single coordinate of a point in raw bytes form.
g1CoordLen, g2CoordLen :: Int
g1CoordLen = 48
g2CoordLen = 96  -- each coordinate is an element of two-dimensional field Fr2

-- | Length of 'Bls12381Fr' in raw bytes form.
frLen :: Int
frLen = 32

-- Meta bits
----------------------------------------------------------------------------

-- | This bit designates whether the point is represented in compressed form
-- (only X coordinate), or uncompressed form (both X and Y coordinates).
compressionBit :: Int
compressionBit = 7

-- | This bit designates whether the given point is at infinity.
--
-- If so, all other bytes should be zeros.
infinityBit :: Int
infinityBit = 6

-- | This bit is set iff "this point is in compressed form /and/ it is not the
-- point at infinity /and/ its y-coordinate is the lexicographically largest of
-- the two associated with the encoded x-coordinate".
--
-- Fortunatelly, this flag seems to be not relevant for us at the moment.
flag3Bit :: Int
flag3Bit = 5

----------------------------------------------------------------------------
-- Other constants
----------------------------------------------------------------------------

-- | @1@ represented in G1 - as the libraries used by Tezos see it.
--
-- Taken from here: <https://github.com/ethereum/py_ecc/blob/3f644b4c07c8270b8fbe989eb799766aca66face/py_ecc/optimized_bls12_381/optimized_curve.py#L34>.
g1One :: Bls12381G1
g1One = Bls12381G1 $ CW.toA $ CW.J
  3685416753713387016781088315183077757961620795782546409894578378688607592378376318836054947676345821548104185464507
  1339506544944476473020471379941921221584933875938349620426543736416511423956333506472724655353366534992391756441569
  1

-- | @1@ represented in G2.
g2One :: Bls12381G2
g2One = Bls12381G2 $ CW.toA $ CW.J
  (GF.toE
    [ 352701069587466618187139116011060144890029952792775240219908644239793785735715026873347600343865175952761926303160
    , 3059144344244213709971259814753781636986470325476647558659373206291635324768958432433509563104347017837885763365758
    ])
  (GF.toE
    [ 1985150602287291935568054521177171638300868978215655730859378665066344726373823718423869104263333984641494340347905
    , 927553665492332455747201965776037880757740193453592970025027978793976877002675564980949289727957565575433344219582
    ])
  (GF.toE [1, 0])