From 6beb6c5f9ab9a48a5b3e92b9a02b01150d89e4a1 Mon Sep 17 00:00:00 2001
From: Quaternions <krakow20@gmail.com>
Date: Thu, 16 Jan 2025 20:36:03 -0800
Subject: [PATCH] implement bot file

---
 lib/snf/src/bot.rs | 363 +++++++++++++++++++++++++++++++++++++--------
 1 file changed, 302 insertions(+), 61 deletions(-)

diff --git a/lib/snf/src/bot.rs b/lib/snf/src/bot.rs
index ffdaee8..3712341 100644
--- a/lib/snf/src/bot.rs
+++ b/lib/snf/src/bot.rs
@@ -1,98 +1,339 @@
-use binrw::{BinReaderExt, binrw};
+use binrw::{binrw,BinReaderExt,BinWrite,BinWriterExt};
+
+use crate::newtypes;
+use crate::file::BlockId;
+use strafesnet_common::physics::Time;
+
+type TimedPhysicsInstruction=strafesnet_common::instruction::TimedInstruction<strafesnet_common::physics::Instruction,strafesnet_common::physics::TimeInner>;
 
 #[derive(Debug)]
 pub enum Error{
-	InvalidHeader,
+	InvalidHeader(binrw::Error),
 	InvalidSegment(binrw::Error),
+	SegmentConvert(newtypes::integer::RatioError),
+	InstructionConvert(newtypes::physics::InstructionConvert),
+	InstructionWrite(binrw::Error),
 	InvalidSegmentId(SegmentId),
+	InvalidData(binrw::Error),
+	IO(std::io::Error),
 	File(crate::file::Error),
 }
 
+// Bot files are simply the sequence of instructions that the physics received during the run.
+// The instructions are partitioned into timestamped blocks for ease of streaming.
+//
+// Keyframe information required for efficient seeking
+// is part of a different file, and is generated from this file.
+
 /* block types
 
 BLOCK_BOT_HEADER:
-u128 map_resource_uuid //which map is this bot running
-//don't include style info in bot header because it's in the simulation state
-//blocks are laid out in chronological order, but indices may jump around.
-u64 num_segments
+// Tegments are laid out in chronological order,
+// but block_id is not necessarily in ascending order.
+//
+// This is to place the final segment close to the start of the file,
+// which allows the duration of the bot to be conveniently calculated
+// from the first and last instruction timestamps.
+u32 num_segments
 for _ in 0..num_segments{
-	i64 time //simulation_state timestamp
-	u64 block_id
+	i64 time
+	u32 instruction_count
+	u32 block_id
 }
 
 BLOCK_BOT_SEGMENT:
-//format version indicates what version of these structures to use
-SimulationState simulation_state //SimulationState is just under ClientState which includes Play/Pause events that the simulation doesn't know about.
-//to read, greedily decode instructions until eof
-loop{
-	//delta encode as much as possible (time,mousepos)
-	//strafe ticks are implied
-	//physics can be implied in an input-only bot file
-	TimedInstruction<SimulationInstruction> instruction
+// segments can potentially be losslessly compressed!
+for _ in 0..instruction_count{
+	// TODO: delta encode as much as possible (time,mousepos)
+	i64 time
+	physics::Instruction instruction
 }
 
 */
 
-//error hiding mock code
-mod simulation{
-	#[super::binrw]
-	#[brw(little)]
-	pub struct State{}
-	#[super::binrw]
-	#[brw(little)]
-	pub struct Instruction{}
+#[binrw]
+#[brw(little)]
+struct SegmentHeader{
+	time:i64,
+	instruction_count:u32,
+	block_id:BlockId,
+}
+#[binrw]
+#[brw(little)]
+struct Header{
+	num_segments:u32,
+	#[br(count=num_segments)]
+	segments:Vec<SegmentHeader>,
 }
-// mod instruction{
-	// #[super::binrw]
-	// #[brw(little)]
-	// pub struct TimedInstruction<Instruction:binrw::BinRead+binrw::BinWrite>{
-		// time:u64,
-		// instruction:Instruction
-	// }
-// }
-// mod timeline{
-	// #[super::binrw]
-	// #[brw(little)]
-	// pub struct Timeline<Instruction:binrw::BinRead+binrw::BinWrite>{
-		// #[bw(try_calc(u32::try_from(instructions.len())))]
-		// instruction_count:u32,
-		// #[br(count=instruction_count)]
-		// instructions:Vec<super::instruction::TimedInstruction<Instruction>>
-	// }
-// }
-
-//serious code
 
 #[binrw]
 #[brw(little)]
 #[derive(Clone,Copy,Debug,id::Id)]
-pub struct SegmentId(u32);
+struct SegmentId(u32);
 
-#[binrw]
-#[brw(little)]
 pub struct Segment{
-	state:simulation::State,
-	//#[bw(try_calc(u32::try_from(instructions.len())))]
-	//instruction_count:u32,
-	//#[br(count=instruction_count)]
-	//instructions:Vec<instruction::TimedInstruction<simulation::Instruction>>
+	pub instructions:Vec<TimedPhysicsInstruction>
+}
 
-	//please remember that strafe ticks are implicit! 33% smaller bot files
+#[derive(Clone,Copy,Debug)]
+pub struct SegmentInfo{
+	/// time of the first instruction in this segment.
+	time:Time,
+	instruction_count:u32,
+	/// How many total instructions in segments up to and including this segment
+	/// Alternatively, the id of the first instruction be in the _next_ segment
+	instructions_subtotal:u64,
+	block_id:BlockId,
 }
 
 pub struct StreamableBot<R:BinReaderExt>{
 	file:crate::file::File<R>,
-	//timeline:timeline::Timeline<SegmentId>,
-	segment_id_to_block_id:Vec<crate::file::BlockId>,
+	segment_map:Vec<SegmentInfo>,
 }
 impl<R:BinReaderExt> StreamableBot<R>{
-	pub(crate) fn new(file:crate::file::File<R>)->Result<Self,Error>{
-		Err(Error::InvalidHeader)
+	pub(crate) fn new(mut file:crate::file::File<R>)->Result<Self,Error>{
+		//assume the file seek is in the right place to start reading header
+		let header:Header=file.data_mut().read_le().map_err(Error::InvalidHeader)?;
+		let mut instructions_subtotal=0;
+		let segment_map=header.segments.into_iter().map(|SegmentHeader{time,instruction_count,block_id}|{
+			instructions_subtotal+=instruction_count as u64;
+			SegmentInfo{
+				time:Time::raw(time),
+				instruction_count,
+				instructions_subtotal,
+				block_id,
+			}
+		}).collect();
+		Ok(Self{
+			file,
+			segment_map,
+		})
 	}
-	pub fn load_segment(&mut self,segment_id:SegmentId)->Result<Segment,Error>{
-		let block_id=*self.segment_id_to_block_id.get(segment_id.get() as usize).ok_or(Error::InvalidSegmentId(segment_id))?;
-		let mut block=self.file.block_reader(block_id).map_err(Error::File)?;
-		let segment=block.read_le().map_err(Error::InvalidSegment)?;
+	fn get_segment_info(&self,segment_id:SegmentId)->Result<SegmentInfo,Error>{
+		Ok(*self.segment_map.get(segment_id.get() as usize).ok_or(Error::InvalidSegmentId(segment_id))?)
+	}
+	pub fn find_segments_instruction_range(&self,start_instruction:u64,end_instruction:u64)->&[SegmentInfo]{
+		let start=self.segment_map.partition_point(|segment_info|segment_info.instructions_subtotal<start_instruction);
+		let end=self.segment_map.partition_point(|segment_info|segment_info.instructions_subtotal<end_instruction);
+		&self.segment_map[start..=end]
+	}
+	// pub fn find_segments_time_range(&self,start_time:Time,end_time:Time)->&[SegmentInfo]{
+	// 	// TODO: This is off by one, both should be one less
+	// 	let start=self.segment_map.partition_point(|segment_info|segment_info.time<start_time);
+	// 	let end=self.segment_map.partition_point(|segment_info|segment_info.time<end_time);
+	// 	&self.segment_map[start..=end]
+	// }
+	fn append_to_segment(&mut self,segment_info:SegmentInfo,segment:&mut Segment)->Result<(),Error>{
+		let mut block=self.file.block_reader(segment_info.block_id).map_err(Error::File)?;
+		for _ in 0..segment_info.instruction_count{
+			let instruction:newtypes::physics::TimedInstruction=block.read_le().map_err(Error::InvalidSegment)?;
+			segment.instructions.push(instruction.try_into().map_err(Error::SegmentConvert)?);
+		}
+		Ok(())
+	}
+	pub fn load_segment(&mut self,segment_info:SegmentInfo)->Result<Segment,Error>{
+		let mut segment=Segment{
+			instructions:Vec::with_capacity(segment_info.instruction_count as usize),
+		};
+		self.append_to_segment(segment_info,&mut segment)?;
+		Ok(segment)
+	}
+	pub fn read_all(&mut self)->Result<Segment,Error>{
+		let mut segment=Segment{
+			instructions:Vec::new(),
+		};
+		for i in 0..self.segment_map.len(){
+			let segment_info=self.segment_map[i];
+			self.append_to_segment(segment_info,&mut segment)?;
+		}
 		Ok(segment)
 	}
 }
+
+const MAX_BLOCK_SIZE:usize=64*1024;//64 kB
+pub fn write_bot<W:BinWriterExt>(mut writer:W,instructions:impl IntoIterator<Item=TimedPhysicsInstruction>)->Result<(),Error>{
+	// decide which instructions to put in which segment
+	// write segment 1 to block 1
+	// write segment N to block 2
+	// write rest of segments
+	// 1 2 3 4 5
+	// becomes
+	// [1 5] 2 3 4
+	struct SegmentHeaderInfo{
+		time:Time,
+		instruction_count:u32,
+		range:core::ops::Range<usize>
+	}
+
+	let mut segment_header_infos=Vec::new();
+	let mut raw_segments=std::io::Cursor::new(Vec::new());
+
+	// block info
+	let mut start_time=Time::ZERO;
+	let mut start_position=raw_segments.position() as usize;
+	let mut instruction_count=0;
+
+	let mut last_position=start_position;
+
+	let mut iter=instructions.into_iter();
+
+	macro_rules! collect_instruction{
+		($instruction:expr)=>{
+			let time=$instruction.time;
+			let instruction_writable:newtypes::physics::TimedInstruction=$instruction.try_into().map_err(Error::InstructionConvert)?;
+			instruction_writable.write_le(&mut raw_segments).map_err(Error::InstructionWrite)?;
+			instruction_count+=1;
+			let position=raw_segments.position() as usize;
+			// exceeds max block size
+			if MAX_BLOCK_SIZE<position-last_position{
+				segment_header_infos.push(SegmentHeaderInfo{
+					time:start_time,
+					instruction_count,
+					range:start_position..last_position,
+				});
+				start_position=last_position;
+				instruction_count=0;
+				start_time=time;
+			}
+			last_position=position;
+		}
+	}
+
+	// unroll one loop iteration to grab the starting time
+	if let Some(instruction)=iter.next(){
+		start_time=instruction.time;
+		collect_instruction!(instruction);
+	}
+
+	for instruction in iter{
+		collect_instruction!(instruction);
+	}
+	//last block, whatever size it happens to be
+	{
+		let final_position=raw_segments.position() as usize;
+		segment_header_infos.push(SegmentHeaderInfo{
+			time:start_time,
+			instruction_count,
+			range:start_position..final_position,
+		});
+	}
+	// drop cursor
+	let raw_segments=raw_segments.into_inner();
+
+	let num_segments=segment_header_infos.len();
+
+	// segments list is in chronological order
+	let make_segment_header=|block_id,&SegmentHeaderInfo{time,instruction_count,range:ref _range}|SegmentHeader{
+		time:time.get(),
+		instruction_count,
+		block_id,
+	};
+	let segments=if 2<num_segments{
+		let mut segments=Vec::with_capacity(num_segments);
+		// segment 1 is second block
+		if let Some(seg)=segment_header_infos.first(){
+			segments.push(make_segment_header(BlockId::new(1),seg));
+		}
+		// rest of segments start at fourth block
+		for (i,seg) in segment_header_infos[1..num_segments-1].iter().enumerate(){
+			make_segment_header(BlockId::new(3+i as u32),seg);
+		}
+		// last segment is third block
+		if let Some(seg)=segment_header_infos.last(){
+			segments.push(make_segment_header(BlockId::new(2),seg));
+		}
+		segments
+	}else{
+		// all segments in order
+		segment_header_infos.iter().enumerate().map(|(i,seg)|
+			make_segment_header(BlockId::new(1+i as u32),seg)
+		).collect()
+	};
+
+	let header=Header{
+		num_segments:num_segments as u32,
+		segments,
+	};
+
+	// map header is +1
+	let block_count=1+num_segments as u32;
+
+	let mut offset=crate::file::Header::calculate_size(block_count) as u64;
+	// block_location is one longer than block_count
+	let mut block_location=Vec::with_capacity(1+block_count as usize);
+
+	//probe header length
+	let mut bot_header_data=Vec::new();
+	binrw::BinWrite::write_le(&header,&mut std::io::Cursor::new(&mut bot_header_data)).map_err(Error::InvalidData)?;
+
+	// the first block location is the map header
+	block_location.push(offset);
+	offset+=bot_header_data.len() as u64;
+	block_location.push(offset);
+
+	// priming includes file header + first 3 blocks [bot header, first segment, last segment]
+	let priming=if 2<num_segments{
+		// segment 1 is block 2
+		if let Some(seg)=segment_header_infos.first(){
+			offset+=seg.range.len() as u64;
+			block_location.push(offset);
+		}
+		// last segment is block 3
+		if let Some(seg)=segment_header_infos.last(){
+			offset+=seg.range.len() as u64;
+			block_location.push(offset);
+		}
+
+		let priming=offset;
+
+		// rest of segments
+		for seg in &segment_header_infos[1..num_segments-1]{
+			offset+=seg.range.len() as u64;
+			block_location.push(offset);
+		}
+		priming
+	}else{
+		// all segments in order
+		for seg in &segment_header_infos{
+			offset+=seg.range.len() as u64;
+			block_location.push(offset);
+		}
+		offset
+	};
+
+	let file_header=crate::file::Header{
+		fourcc:crate::file::FourCC::Bot,
+		version:0,
+		priming,
+		resource:0,
+		block_count,
+		block_location,
+	};
+
+	// write file header
+	writer.write_le(&file_header).map_err(Error::InvalidData)?;
+	// write bot header
+	writer.write(&bot_header_data).map_err(Error::IO)?;
+
+	// write blocks
+	if 2<num_segments{
+		// segment 1 is block 2
+		if let Some(seg)=segment_header_infos.first(){
+			writer.write(&raw_segments[seg.range.clone()]).map_err(Error::IO)?;
+		}
+		// last segment is block 3
+		if let Some(seg)=segment_header_infos.last(){
+			writer.write(&raw_segments[seg.range.clone()]).map_err(Error::IO)?;
+		}
+		// rest of segments
+		for seg in &segment_header_infos[1..num_segments-1]{
+			writer.write(&raw_segments[seg.range.clone()]).map_err(Error::IO)?;
+		}
+	}else{
+		// all segments in order
+		for seg in segment_header_infos{
+			writer.write(&raw_segments[seg.range]).map_err(Error::IO)?;
+		}
+	}
+	Ok(())
+}