corevm 0.1.22-rc.3

use crate::{JamLogger, JamOuterVm};
use alloc::vec::Vec;
use codec::{Decode, DecodeAll};
use corevm_engine::{initial_state_hash, Engine};
use corevm_host::{CoreVmInstruction, CoreVmOutput, Range, RangeSet, StorageKey, VmSpec};
use jam_pvm_common::accumulate::{get, remove, set};
use jam_types::{
	AccumulateItem, CodeHash, Encode, Hash, RefineContext, ServiceId, Slot, TransferRecord,
	WorkError, WorkOutput, WorkPackageHash, WorkPayload,
};
use log::debug;

#[allow(dead_code)]
struct Service;
jam_pvm_common::declare_service!(Service);

impl jam_pvm_common::Service for Service {
	fn refine(
		_id: ServiceId,
		payload: WorkPayload,
		_package_hash: WorkPackageHash,
		_context: RefineContext,
		_auth_code_hash: CodeHash,
	) -> WorkOutput {
		JamLogger::init();
		let payload = DecodeAll::decode_all(&mut &payload[..]).expect("Failed to decode payload");
		let engine = Engine::new(payload, JamOuterVm).expect("Failed to create engine");
		let (output, _outer_vm) = engine.run().expect("Failed to execute the code");
		WorkOutput(output.encode())
	}

	fn accumulate(_slot: Slot, id: ServiceId, items: Vec<AccumulateItem>) -> Option<Hash> {
		JamLogger::init();
		debug!("Running accumulate");
		for item in items.into_iter() {
			if let Err(e) = update_pages_in_storage(item) {
				debug!("Failed to update pages in the storage: {e:?}");
			} else {
				debug!("Updated pages in the storage of {id:#x}");
			}
		}
		None
	}

	fn on_transfer(_slot: Slot, _id: ServiceId, items: Vec<TransferRecord>) {
		JamLogger::init();
		for item in items.into_iter() {
			if let Ok(instr) = CoreVmInstruction::decode(&mut &item.memo[..]) {
				match instr {
					CoreVmInstruction::SetCode { code_hash, host } => {
						if let Some(owner) = get::<ServiceId>(StorageKey::Owner) {
							if owner != item.source {
								debug!("Not the owner");
								break;
							}
						} else if set(StorageKey::Owner, item.source).is_err() {
							debug!("Failed to set owner");
						}
						debug!("Set code hash {code_hash} for {host:#x}");
						if set(StorageKey::GuestCode, (code_hash, host)).is_err() {
							debug!("Failed to set code hash");
						}
						// Now clear memory pages.
						if let Some(spec) = get::<VmSpec>(StorageKey::VmSpec) {
							for range in spec.state.resident_pages.as_slice().iter() {
								for address in range.start..range.end {
									remove(StorageKey::PageHash(address));
								}
							}
						}
						// Clear the VM state.
						remove(StorageKey::StateHash);
						remove(StorageKey::VmSpec);
						remove(StorageKey::VideoMode);
					},
					CoreVmInstruction::SetOwner(owner) => {
						if let Some(owner) = get::<ServiceId>(StorageKey::Owner) {
							if owner != item.source {
								debug!("Not the owner");
								break;
							}
						}
						if set(StorageKey::Owner, owner).is_err() {
							debug!("Failed to set owner");
						}
					},
				}
			}
		}
	}
}

fn update_pages_in_storage(item: AccumulateItem) -> Result<(), AccumulateError> {
	use AccumulateError::*;
	let output = item.result?;
	let output = CoreVmOutput::decode_all(&mut &output[..])?;

	// Check the code hash, storing it if not already set.
	match get::<(Hash, ServiceId)>(StorageKey::GuestCode) {
		Some((guest_code_hash, _host_id)) =>
			if guest_code_hash != output.guest_code_hash {
				return Err(GuestCodeMismatch);
			},
		None => {
			// Store the guest code hash and host ID.
			set(StorageKey::GuestCode, (output.guest_code_hash, output.guest_code_host))?;
		},
	};

	// Update VM state hash.
	match get::<Hash>(StorageKey::StateHash) {
		Some(their_hash) =>
			if their_hash != output.old_hash {
				return Err(InvalidPriorState);
			},
		None => {
			// No prior state, but may be we started from the initial state?
			let initial = initial_state_hash();
			if initial != output.old_hash {
				return Err(NoPriorState);
			}
		},
	}
	set(StorageKey::StateHash, output.new_hash)?;
	let mut resident_pages = get::<RangeSet>(StorageKey::ResidentPages).unwrap_or_default();
	// Check imported pages' hashes.
	for (address, our_hash) in output.touched_imported_pages.iter() {
		let key = StorageKey::PageHash(*address);
		let their_hash: Hash = get(&key).ok_or(NoSuchPage)?;
		if &their_hash != our_hash {
			return Err(PageHashMismatch);
		}
	}
	// Update memory pages hashes.
	for (address, hash) in output.updated_pages.iter() {
		let key = StorageKey::PageHash(*address);
		if hash == &[0; 32] {
			// The page was freed.
			get::<Hash>(&key).ok_or(NoSuchPage)?;
			resident_pages.remove(&Range::new(*address, address + 1));
			remove(&key);
		} else {
			resident_pages.insert(Range::new(*address, address + 1));
			set(&key, hash)?;
		}
	}
	// Update video mode.
	if let Some(mode) = output.video {
		set(StorageKey::VideoMode, mode)?;
	}
	// Store VM output specification.
	let spec = VmSpec {
		exports_root: item.exports_root,
		output: output.vm_output,
		state: output.vm_state,
	};
	set(StorageKey::VmSpec, spec)?;

	Ok(())
}

#[derive(Debug)]
#[allow(dead_code)]
enum AccumulateError {
	Work(WorkError),
	Codec(codec::Error),
	Api(jam_pvm_common::ApiError),
	NoPriorState,
	InvalidPriorState,
	NoSuchPage,
	PageHashMismatch,
	GuestCodeMismatch,
}

impl From<WorkError> for AccumulateError {
	fn from(e: WorkError) -> Self {
		Self::Work(e)
	}
}

impl From<codec::Error> for AccumulateError {
	fn from(e: codec::Error) -> Self {
		Self::Codec(e)
	}
}

impl From<jam_pvm_common::ApiError> for AccumulateError {
	fn from(e: jam_pvm_common::ApiError) -> Self {
		Self::Api(e)
	}
}