type class combiner with polymorphic list type/instance?

Stream: lean4

[David Stainton 🦡]

New to Lean. Trying to keep my question concise.
I need some guidance or advice on using Lean.
I should say that I'm only about half-way through reading Functional Programming in Lean.
So... perhaps you can tell me, go and read these chapters over there and then revisit this problem.

I wrote a KEM type class.
My goal is to write a KEM combiner which combines an arbitrary number of KEMs together.
Seems like I can't just put this into the Combiner type class definition...

List (Σ α : Type, KEM α × α)

Perhaps my KEM type class is not good and I need something else.
Before I realized exactly how I wanted to use Lean, I was curious how Lean
might be used to create symbolic cryptographic models, so @Mario Carneiro
wrote a really impressive KEM definition, complete with proofs of "KEM correctness".
My intention is to revisit that later after I learn Lean a bit better.
Here's what my KEM type class looks like... using my beginner Lean functional programming skills:

class Key (key : Type) where
  encode : key → ByteArray
  decode : ByteArray → Option key

class PrivateKey (privkey : Type) extends Key privkey

class PublicKey (pubkey : Type) extends Key pubkey

class KEM (scheme : Type) where
  PublicKeyType : Type
  PrivateKeyType : Type

  name : String
  generateKeyPair : IO (PublicKeyType × PrivateKeyType)
  encapsulate : scheme → PublicKeyType → IO (ByteArray × ByteArray)
  decapsulate : scheme → PrivateKeyType → ByteArray → ByteArray
  privateKeySize : Nat
  publicKeySize : Nat
  encodePrivateKey : PrivateKeyType → ByteArray
  decodePrivateKey : ByteArray → Option PrivateKeyType
  encodePublicKey : PublicKeyType → ByteArray
  decodePublicKey : ByteArray → Option PublicKeyType

I construct my List of KEM like this:

def defaultHash := fun x => Sha256.hash x

instance : Adapter X25519Scheme where
  hash := defaultHash

instance : Adapter X448Scheme where
  hash := defaultHash

instance : Adapter X41417Scheme where
  hash := defaultHash

def X25519AsKEM : KEM X25519Scheme := inferInstance
def X448AsKEM : KEM X448Scheme := inferInstance
def X41417AsKEM : KEM X41417Scheme := inferInstance

def X25519Instance : X25519Scheme := {}
def X448Instance : X448Scheme := {}
def X41417Instance : X41417Scheme := {}

def Schemes : List (Σ α : Type, KEM α × α) :=
  [
    ⟨X25519Scheme, (X25519AsKEM, X25519Instance)⟩,
    ⟨X448Scheme, (X448AsKEM, X448Instance)⟩,
    ⟨X41417Scheme, (X41417AsKEM, X41417Instance)⟩,
  ]

I test the list of KEMs like this:

def testKEM {α : Type} [kemInstance : KEM α] (kem : α) : IO Unit := do
  let (alicePublicKey, alicePrivateKey) ← kemInstance.generateKeyPair
  let (ciphertext, bobSharedSecret) ← kemInstance.encapsulate kem alicePublicKey
  let aliceSharedSecret := kemInstance.decapsulate kem alicePrivateKey ciphertext
  if bobSharedSecret == aliceSharedSecret then
    IO.println s!"KEM test for {kemInstance.name} PASSED."
  else
    panic! "KEM test failed!"

def testAllKEMs : List (Σ α : Type, KEM α × α) → IO Unit
| [] => IO.println "All KEM tests passed!"
| ⟨_, _, kem⟩ :: rest => do
  testKEM kem
  testAllKEMs rest

[Yakov Pechersky]

May I ask, why do you use class and instance, as opposed to structure and def?

[David Stainton 🦡]

Not sure. Would it be better to use structures to represent the KEM instead?

[Yakov Pechersky]

Your testKem is expecting to find the Kem functions using typeclass search. But in your testAllKEMs, you aren't supplying the unwrapped foo : KEM kem back to the TC system

[Yakov Pechersky]

You could change your inductive step to have

def testAllKEMs : List (Σ α : Type, KEM α × α) → IO Unit
| [] => IO.println "All KEM tests passed!"
| ⟨_, inst, kem⟩ :: rest => do
  have := inst
  testKEM kem
  testAllKEMs rest

[David Stainton 🦡]

Okay thanks I'll try changing it to a polymorphic structure. It sounds like that could make it easier to write a Kem combiner. The structure could have a field of the List type above.

[Yakov Pechersky]

Did my example work even with the current class based style?

[David Stainton 🦡]

yeah your change works too.