jam-program-blob-common 0.1.26

use alloc::{borrow::Cow, string::String, vec::Vec};
use codec::{Compact, CompactLen, Decode, Encode};

/// Information on a crate, useful for building conventional metadata of type 0.
#[derive(Clone, PartialEq, Eq, Encode, Decode, Debug)]
pub struct CrateInfo {
	pub name: String,
	pub version: String,
	pub license: String,
	pub authors: Vec<String>,
}

impl core::fmt::Display for CrateInfo {
	fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
		let name = match self.name.trim() {
			"" => "<no-name>",
			other => other,
		};
		let version = match self.version.trim() {
			"" => "<no-version>",
			other => other,
		};
		write!(f, "{name} v{version} by ")?;
		if self.authors.is_empty() || self.authors.iter().all(|s| s.trim().is_empty()) {
			write!(f, "<no-authors>")?;
		} else {
			let mut iter = self.authors.iter().map(|s| s.trim()).filter(|s| !s.is_empty());
			if let Some(first) = iter.next() {
				write!(f, "{first}")?;
			}
			for author in iter {
				write!(f, ", {author}")?;
			}
		}
		Ok(())
	}
}

/// Information which, when encoded, could fill a program blob's metadata.
#[derive(Clone, PartialEq, Eq, Encode, Decode, Debug)]
pub enum ConventionalMetadata {
	Info(CrateInfo),
}

/// Encode and write `ConventionalMetadata` to the `output`.
pub fn write_metadata(
	metadata: &ConventionalMetadata,
	output: &mut Vec<u8>,
) -> Result<(), &'static str> {
	let len = metadata.encoded_size();
	write_var(u32::try_from(len).map_err(|_| "metadata too large")?, output);
	metadata.encode_to(output);
	Ok(())
}

/// Read and decode `ConventionalMetadata`.
pub fn read_metadata(bytes: &mut &[u8]) -> Result<ConventionalMetadata, codec::Error> {
	let mut metadata = read_metadata_slice(bytes).ok_or("Failed to parse metadata")?;
	ConventionalMetadata::decode(&mut metadata)
}

/// Read `ConventionalMetadata` as bytes.
pub fn read_metadata_slice<'a>(bytes: &mut &'a [u8]) -> Option<&'a [u8]> {
	let offset = read_var(bytes)?;
	let metadata = read_slice(bytes, offset)?;
	Some(metadata)
}

/// A JAM-specific program blob.
pub struct ProgramBlob<'a> {
	pub metadata: Cow<'a, [u8]>,
	pub ro_data: Cow<'a, [u8]>,
	pub rw_data: Cow<'a, [u8]>,
	pub code_blob: Cow<'a, [u8]>,
	pub rw_data_padding_pages: u16,
	pub stack_size: u32,
}

fn read_u24(bytes: &mut &[u8]) -> Option<u32> {
	let xs = bytes.get(..3)?;
	*bytes = &bytes[3..];
	Some(u32::from_le_bytes([xs[0], xs[1], xs[2], 0]))
}

fn write_u24(value: u32, output: &mut Vec<u8>) -> Result<(), ()> {
	if value >= (1 << 24) {
		return Err(());
	}

	output.extend_from_slice(&value.to_le_bytes()[0..3]);
	Ok(())
}

fn read_u16(bytes: &mut &[u8]) -> Option<u16> {
	let xs = bytes.get(..2)?;
	*bytes = &bytes[2..];
	Some(u16::from_le_bytes([xs[0], xs[1]]))
}

fn read_u32(bytes: &mut &[u8]) -> Option<u32> {
	let xs = bytes.get(..4)?;
	*bytes = &bytes[4..];
	Some(u32::from_le_bytes([xs[0], xs[1], xs[2], xs[3]]))
}

fn read_var(bytes: &mut &[u8]) -> Option<u32> {
	Some(Compact::<u32>::decode(bytes).ok()?.0)
}

fn write_var(value: u32, output: &mut Vec<u8>) {
	Compact::<u32>(value).encode_to(output)
}

fn read_cow<'a>(bytes: &mut &'a [u8], length: u32) -> Option<Cow<'a, [u8]>> {
	read_slice(bytes, length).map(Cow::Borrowed)
}

fn read_slice<'a>(bytes: &mut &'a [u8], length: u32) -> Option<&'a [u8]> {
	let length = length as usize;
	let slice = bytes.get(..length)?;
	*bytes = &bytes[length..];
	Some(slice)
}

impl<'a> ProgramBlob<'a> {
	pub fn from_bytes(mut bytes: &'a [u8]) -> Option<Self> {
		let metadata = Cow::Borrowed(read_metadata_slice(&mut bytes)?);
		let ro_data_len = read_u24(&mut bytes)?;
		let rw_data_len = read_u24(&mut bytes)?;
		let rw_data_padding_pages = read_u16(&mut bytes)?;
		let stack_size = read_u24(&mut bytes)?;
		let ro_data = read_cow(&mut bytes, ro_data_len)?;
		let rw_data = read_cow(&mut bytes, rw_data_len)?;
		let code_blob_len = read_u32(&mut bytes)?;
		let code_blob = read_cow(&mut bytes, code_blob_len)?;

		if !bytes.is_empty() {
			return None;
		}

		Some(ProgramBlob {
			metadata,
			rw_data_padding_pages,
			stack_size,
			ro_data,
			rw_data,
			code_blob,
		})
	}

	pub fn to_vec(&self) -> Result<Vec<u8>, &'static str> {
		let mut output = Vec::new();
		write_var(
			u32::try_from(self.metadata.len()).map_err(|_| "metadata too large")?,
			&mut output,
		);
		output.extend_from_slice(&self.metadata);
		write_u24(u32::try_from(self.ro_data.len()).map_err(|_| "too large RO data")?, &mut output)
			.map_err(|_| "too large RO data")?;
		write_u24(u32::try_from(self.rw_data.len()).map_err(|_| "too large RW data")?, &mut output)
			.map_err(|_| "too large RW data")?;
		output.extend_from_slice(&self.rw_data_padding_pages.to_le_bytes());
		write_u24(self.stack_size, &mut output).map_err(|_| "too large stack size")?;
		output.extend_from_slice(&self.ro_data);
		output.extend_from_slice(&self.rw_data);
		output.extend_from_slice(
			&u32::try_from(self.code_blob.len()).map_err(|_| "too large code")?.to_le_bytes(),
		);
		output.extend_from_slice(&self.code_blob);
		Ok(output)
	}
}

#[cfg(feature = "polkavm")]
impl From<ProgramBlob<'_>> for polkavm::ProgramParts {
	fn from(other: ProgramBlob<'_>) -> Self {
		let mut parts = polkavm::ProgramParts::default();
		parts.ro_data_size = other.ro_data.len() as u32;
		parts.rw_data_size = other.rw_data.len().next_multiple_of(4096) as u32 +
			other.rw_data_padding_pages as u32 * 4096;
		parts.stack_size = other.stack_size;
		parts.ro_data = other.ro_data.into();
		parts.rw_data = other.rw_data.into();
		parts.code_and_jump_table = other.code_blob.into();
		parts.is_64_bit = true;
		parts
	}
}

#[cfg(feature = "polkavm")]
impl<'a> ProgramBlob<'a> {
	pub fn from_pvm(parts: &'a polkavm::ProgramParts, metadata: Cow<'a, [u8]>) -> Self {
		// Pad RO section with zeroes.
		let mut ro_data = parts.ro_data.to_vec();
		ro_data.resize(parts.ro_data_size as usize, 0);
		// Calculate the padding for RW section.
		let padding = (parts.rw_data_size as usize).next_multiple_of(4096) -
			parts.rw_data.len().next_multiple_of(4096);
		let rw_data_padding_pages = padding / 4096;
		let rw_data_padding_pages =
			rw_data_padding_pages.try_into().expect("The RW data section is too big");
		Self {
			metadata,
			ro_data: ro_data.into(),
			rw_data: (&parts.rw_data[..]).into(),
			code_blob: (&parts.code_and_jump_table[..]).into(),
			rw_data_padding_pages,
			stack_size: parts.stack_size,
		}
	}
}

/// A CoreVM-specific program blob.
pub struct CoreVmProgramBlob<'a, 'b> {
	/// Serialized conventional metadata.
	pub metadata: Cow<'a, [u8]>,
	/// PVM binary blob.
	pub pvm_blob: Cow<'b, [u8]>,
}

impl<'a> CoreVmProgramBlob<'a, 'a> {
	/// Deserialize from bytes.
	pub fn from_bytes(mut bytes: &'a [u8]) -> Option<Self> {
		let metadata = Cow::Borrowed(read_metadata_slice(&mut bytes)?);
		let pvm_blob = Cow::Borrowed(bytes);
		Some(Self { metadata, pvm_blob })
	}

	/// Serialize into bytes.
	pub fn to_vec(&self) -> Result<Vec<u8>, &'static str> {
		let metadata_len = u32::try_from(self.metadata.len()).map_err(|_| "metadata too large")?;
		let mut output = Vec::with_capacity(
			Compact::<u32>::compact_len(&metadata_len) +
				metadata_len as usize +
				self.pvm_blob.len(),
		);
		write_var(metadata_len, &mut output);
		output.extend_from_slice(&self.metadata);
		output.extend_from_slice(&self.pvm_blob);
		Ok(output)
	}
}

#[cfg(test)]
mod tests {
	use super::*;

	#[test]
	fn read_write_u24() {
		let mut output = Vec::new();
		write_u24(0x00345678, &mut output).unwrap();
		assert_eq!(read_u24(&mut &output[..]), Some(0x00345678));

		assert!(write_u24(0x00ffffff, &mut output).is_ok());
		assert!(write_u24(0x01000000, &mut output).is_err());
	}

	#[test]
	fn read_write_var() {
		let mut output = Vec::new();
		let vals = [0x00345678, 0x00, 0x01, 0x7f, 0x80, 0xffffffff];
		for i in vals.into_iter() {
			write_var(i, &mut output);
		}
		let mut cursor = output.as_ref();
		for i in vals.into_iter() {
			assert_eq!(read_var(&mut cursor), Some(i));
		}
	}

	#[test]
	fn metadata_read_write() {
		let info = CrateInfo {
			name: "x".into(),
			version: "y".into(),
			license: "z".into(),
			authors: vec!["w".into(), "v".into()],
		};
		let metadata = ConventionalMetadata::Info(info);
		let mut output = Vec::new();
		write_metadata(&metadata, &mut output).unwrap();
		let actual = read_metadata(&mut &output[..]).unwrap();
		assert_eq!(metadata, actual);
	}

	#[test]
	fn corevm_program_blob_read_write() {
		let info = CrateInfo {
			name: "x".into(),
			version: "y".into(),
			license: "z".into(),
			authors: vec!["w".into(), "v".into()],
		};
		let metadata = ConventionalMetadata::Info(info).encode().into();
		// This is just a program composed of a single return.
		let pvm_blob = vec![0, 0, 2, 50, 0, 1];
		let corevm_blob = CoreVmProgramBlob { metadata, pvm_blob: pvm_blob.as_slice().into() };
		let corevm_blob_bytes = corevm_blob.to_vec().unwrap();
		let actual_blob = CoreVmProgramBlob::from_bytes(&corevm_blob_bytes[..]).unwrap();
		assert_eq!(corevm_blob.metadata, actual_blob.metadata);
		assert_eq!(corevm_blob.pvm_blob, actual_blob.pvm_blob);
		assert_eq!(pvm_blob.as_slice(), actual_blob.pvm_blob.as_ref());
	}
}