Collapse extension enum into block

6 files changed, 146 insertions, 112 deletions

diff --git a/gifed/src/block/extension/graphiccontrol.rs b/gifed/src/block/extension/graphiccontrol.rs
index b595554..8b62160 100644
--- a/gifed/src/block/extension/graphiccontrol.rs
+++ b/gifed/src/block/extension/graphiccontrol.rs
@@ -22,7 +22,7 @@ impl GraphicControl {
 		};
 
 		ret.set_disposal_method(disposal_method);
-		ret.user_input(user_input_flag);
+		ret.set_user_input(user_input_flag);
 		ret.transparency(transparency_flag);
 
 		ret
@@ -51,7 +51,11 @@ impl GraphicControl {
 		self.transparency_index
 	}
 
-	pub fn user_input(&mut self, flag: bool) {
+	pub fn user_input(&self) -> bool {
+		self.packed & 0b000_000_1_0 > 0
+	}
+
+	pub fn set_user_input(&mut self, flag: bool) {
 		if flag {
 			self.packed |= 0b000_000_1_0;
 		} else {
diff --git a/gifed/src/block/extension/mod.rs b/gifed/src/block/extension/mod.rs
index fb5eb20..99ac88b 100644
--- a/gifed/src/block/extension/mod.rs
+++ b/gifed/src/block/extension/mod.rs
@@ -1,49 +1,5 @@
 mod application;
 mod graphiccontrol;
 
-pub use graphiccontrol::{DisposalMethod, GraphicControl};
-
 pub use self::application::Application;
-
-pub enum Extension {
-	GraphicControl(GraphicControl),
-	Looping(u16),
-	Comment(Vec<u8>), // Plain Text
-	Application(Application),
-}
-
-impl From<&Extension> for Box<[u8]> {
-	fn from(ext: &Extension) -> Self {
-		let mut vec = vec![];
-		vec.push(0x21); // Push the extension introducer
-
-		match ext {
-			Extension::GraphicControl(gc) => {
-				vec.push(0xF9); // Graphic control label
-				vec.push(0x04); // Block size for this extension is always 4
-				vec.push(gc.packed);
-				vec.extend_from_slice(&gc.delay_time.to_le_bytes());
-				vec.push(gc.transparency_index);
-			}
-			Extension::Looping(count) => {
-				vec.push(0xFF); // Application extension label
-				vec.push(0x0B); // 11 bytes in this block
-				vec.extend_from_slice(b"NETSCAPE2.0"); // App. ident. and "auth code"
-				vec.push(0x03); // Sub-block length
-				vec.push(0x01); // Identifies netscape looping extension
-				vec.extend_from_slice(&count.to_le_bytes());
-			}
-			Extension::Comment(_) => todo!(),
-			Extension::Application(_) => todo!(),
-		}
-
-		vec.push(0x00); // Zero-length data block indicates end of extension
-		vec.into_boxed_slice()
-	}
-}
-
-impl From<GraphicControl> for Extension {
-	fn from(gce: GraphicControl) -> Self {
-		Extension::GraphicControl(gce)
-	}
-}
+pub use graphiccontrol::{DisposalMethod, GraphicControl};
diff --git a/gifed/src/block/mod.rs b/gifed/src/block/mod.rs
index e35224b..a101600 100644
--- a/gifed/src/block/mod.rs
+++ b/gifed/src/block/mod.rs
@@ -14,12 +14,68 @@ pub use version::Version;
 
 use crate::writer::ImageBuilder;
 
+use self::extension::Application;
+use self::extension::GraphicControl;
+
 pub enum Block {
 	IndexedImage(IndexedImage),
-	Extension(extension::Extension),
+	//TODO: Extension(Extension),
+	GraphicControlExtension(GraphicControl),
+	CommentExtension(Vec<u8>),
+	//TODO: PlainTextExtension(PlainTextExtension),
+	ApplicationExtension(Application),
+	LoopingExtension(LoopCount),
 }
 
 enum WriteBlock {
 	ImageBuilder(ImageBuilder),
 	Block(Block),
 }
+
+pub enum LoopCount {
+	Forever,
+	Number(u16),
+}
+
+pub(crate) fn encode_block(mcs: u8, block: &Block) -> Box<[u8]> {
+	match block {
+		Block::IndexedImage(img) => img.as_boxed_slice(mcs),
+		Block::GraphicControlExtension(_) => encode_extension(block),
+		Block::CommentExtension(_) => encode_extension(block),
+		Block::ApplicationExtension(_) => encode_extension(block),
+		Block::LoopingExtension(_) => encode_extension(block),
+	}
+}
+
+fn encode_extension(block: &Block) -> Box<[u8]> {
+	let mut vec = vec![];
+	vec.push(0x21); // Extension Introducer
+
+	match block {
+		Block::IndexedImage(_) => unreachable!(),
+		Block::GraphicControlExtension(gce) => {
+			vec.push(0xF9); // Graphic control label
+			vec.push(0x04); // Block size for this extension is always 4
+			vec.push(gce.packed);
+			vec.extend_from_slice(&gce.delay_time.to_le_bytes());
+			vec.push(gce.transparency_index);
+		}
+		Block::CommentExtension(comment) => todo!(),
+		Block::ApplicationExtension(app) => todo!(),
+		Block::LoopingExtension(lc) => {
+			vec.push(0xFF); // Application extension label
+			vec.push(0x0B); // 11 bytes in this block
+			vec.extend_from_slice(b"NETSCAPE2.0"); // App. ident. and "auth code"
+			vec.push(0x03); // Sub-block length
+			vec.push(0x01); // Identifies netscape looping extension
+
+			match lc {
+				LoopCount::Forever => vec.extend_from_slice(&[0x00, 0x00]),
+				LoopCount::Number(count) => vec.extend_from_slice(&count.to_le_bytes()),
+			}
+		}
+	}
+
+	vec.push(0x00); // Zero length sub-block indicates end of extension
+	vec.into_boxed_slice()
+}
diff --git a/gifed/src/gif.rs b/gifed/src/gif.rs
index 89aaa64..5f1bd7a 100644
--- a/gifed/src/gif.rs
+++ b/gifed/src/gif.rs
@@ -1,10 +1,12 @@
-use std::{fs::File, io::Write, path::Path};
+use std::{fs::File, io::Write, iter::Peekable, path::Path};
 
 use crate::{
-	block::{extension::Extension, Block, ColorTable, ScreenDescriptor, Version},
+	block::{
+		encode_block, extension::GraphicControl, Block, ColorTable, ScreenDescriptor, Version,
+	},
 	colorimage,
 	writer::GifBuilder,
-	ColorImage,
+	Color, ColorImage,
 };
 pub struct Gif {
 	pub header: Version,
@@ -38,17 +40,8 @@ impl Gif {
 		};
 
 		for block in self.blocks.iter() {
-			match block {
-				Block::IndexedImage(image) => {
-					boxed = image.as_boxed_slice(mcs);
-					out.extend_from_slice(&*boxed);
-				}
-				//Block::BlockedImage(_) => unreachable!(),
-				Block::Extension(ext) => {
-					boxed = ext.into();
-					out.extend_from_slice(&*boxed);
-				}
-			}
+			boxed = encode_block(mcs, block);
+			out.extend_from_slice(&*boxed);
 		}
 
 		// Write Trailer
@@ -84,7 +77,7 @@ impl<'a> Iterator for ImageIterator<'a> {
 			match self.veciter.next() {
 				Some(block) => match block {
 					Block::IndexedImage(img) => break img,
-					Block::Extension(Extension::GraphicControl(gce)) => {
+					Block::GraphicControlExtension(gce) => {
 						if gce.is_transparent() {
 							transparent = Some(gce.transparency_index());
 						} else {
@@ -121,56 +114,84 @@ impl<'a> Iterator for ImageIterator<'a> {
 	}
 }
 
-pub struct FrameIterator<'a> {
-	gif: &'a Gif,
-	veciter: std::slice::Iter<'a, Block>,
-	buffer: Vec<u8>,
-}
-
 pub struct Image<'a> {
-	pub(crate) width: u16,
-	pub(crate) height: u16,
-	left_offset: u16,
-	top_offset: u16,
-	pub(crate) palette: &'a ColorTable,
-	pub(crate) transparent_index: Option<u8>,
-	pub(crate) indicies: &'a [u8],
+	pub width: u16,
+	pub height: u16,
+	pub left_offset: u16,
+	pub top_offset: u16,
+	pub palette: &'a ColorTable,
+	pub transparent_index: Option<u8>,
+	pub indicies: &'a [u8],
 }
 
 impl<'a> Image<'a> {
-	pub fn width(&self) -> u16 {
-		self.width
-	}
+	pub fn rgba(&self) -> Option<Vec<u8>> {
+		let mut rgba = vec![0; self.indicies.len() * 4];
+
+		for (image_index, &color_index) in self.indicies.iter().enumerate() {
+			match self.transparent_index {
+				Some(trans) if trans == color_index => {
+					rgba[image_index as usize * 4] = 0;
+					rgba[image_index * 4 + 1] = 0;
+					rgba[image_index * 4 + 2] = 0;
+					rgba[image_index * 4 + 3] = 0;
+				}
+				_ => {
+					if let Some(color) = self.palette.get(color_index) {
+						rgba[image_index * 4] = color.r;
+						rgba[image_index * 4 + 1] = color.g;
+						rgba[image_index * 4 + 2] = color.b;
+						rgba[image_index * 4 + 3] = 255;
+					} else {
+						return None;
+					}
+				}
+			}
+		}
 
-	pub fn height(&self) -> u16 {
-		self.height
+		Some(rgba)
 	}
 
-	pub fn position(&self) -> (u16, u16) {
-		(self.left_offset, self.top_offset)
-	}
+	pub fn rgb(&self, transparent_replace: Color) -> Option<Vec<u8>> {
+		let mut rgb = vec![0; self.indicies.len() * 3];
 
-	pub fn palette(&self) -> &ColorTable {
-		self.palette
-	}
+		for (image_index, &color_index) in self.indicies.iter().enumerate() {
+			match self.transparent_index {
+				Some(trans) if trans == color_index => {
+					rgb[image_index as usize * 4] = transparent_replace.r;
+					rgb[image_index * 3 + 1] = transparent_replace.g;
+					rgb[image_index * 3 + 2] = transparent_replace.b;
+				}
+				_ => {
+					if let Some(color) = self.palette.get(color_index) {
+						rgb[image_index * 3] = color.r;
+						rgb[image_index * 3 + 1] = color.g;
+						rgb[image_index * 3 + 2] = color.b;
+					} else {
+						return None;
+					}
+				}
+			}
+		}
 
-	pub fn transparent_index(&self) -> Option<u8> {
-		self.transparent_index
+		Some(rgb)
 	}
+}
 
-	pub fn indicies(&self) -> &[u8] {
-		self.indicies
-	}
+pub struct FrameIterator<'a> {
+	gif: &'a Gif,
+	veciter: std::slice::Iter<'a, Block>,
+	buffer: Vec<u8>,
 }
 
 pub struct Frame {
-	width: u16,
-	height: u16,
-	palette: ColorTable,
-	transparent_index: Option<u8>,
-	indicies: Vec<u8>,
-	delay_after_draw: u16,
-	user_input_flag: bool,
+	pub width: u16,
+	pub height: u16,
+	pub palette: ColorTable,
+	pub transparent_index: Option<u8>,
+	pub indicies: Vec<u8>,
+	pub delay_after_draw: u16,
+	pub user_input_flag: bool,
 }
 
 #[cfg(test)]
diff --git a/gifed/src/reader/mod.rs b/gifed/src/reader/mod.rs
index 41494df..966ebde 100644
--- a/gifed/src/reader/mod.rs
+++ b/gifed/src/reader/mod.rs
@@ -10,7 +10,7 @@ use std::{
 
 use crate::{
 	block::{
-		extension::{Application, Extension, GraphicControl},
+		extension::{Application, GraphicControl},
 		Block, ColorTable, CompressedImage, ImageDescriptor, IndexedImage, ScreenDescriptor,
 		Version,
 	},
@@ -104,13 +104,11 @@ impl GifReader {
 					.ok_or(DecodingError::UnexpectedEof)?;
 				reader.skip(1); // Skip block terminator
 
-				Ok(Block::Extension(Extension::GraphicControl(
-					GraphicControl::from(data),
-				)))
+				Ok(Block::GraphicControlExtension(GraphicControl::from(data)))
 			}
-			0xFE => Ok(Block::Extension(Extension::Comment(
+			0xFE => Ok(Block::CommentExtension(
 				reader.take_and_collapse_subblocks(),
-			))),
+			)),
 			0x01 => todo!(), //TODO: do; plain text extension
 			0xFF => {
 				//TODO: error instead of unwraps
@@ -120,11 +118,11 @@ impl GifReader {
 					TryInto::try_into(reader.take(3).unwrap()).unwrap();
 				let data = reader.take_and_collapse_subblocks();
 
-				Ok(Block::Extension(Extension::Application(Application {
+				Ok(Block::ApplicationExtension(Application {
 					identifier,
 					authentication_code,
 					data,
-				})))
+				}))
 			}
 			_ => panic!("Unknown Extension Identifier!"),
 		}
diff --git a/gifed/src/writer/gifbuilder.rs b/gifed/src/writer/gifbuilder.rs
index 57a62e3..337e404 100644
--- a/gifed/src/writer/gifbuilder.rs
+++ b/gifed/src/writer/gifbuilder.rs
@@ -1,6 +1,6 @@
 use std::convert::TryInto;
 
-use crate::block::{extension::Extension, Block, ColorTable, ScreenDescriptor, Version};
+use crate::block::{Block, ColorTable, LoopCount, ScreenDescriptor, Version};
 use crate::writer::ImageBuilder;
 use crate::{EncodingError, Gif};
 
@@ -55,7 +55,7 @@ impl GifBuilder {
 		}
 
 		if let Some(gce) = ib.get_graphic_control() {
-			self.blocks.push(Block::Extension(gce.into()));
+			self.blocks.push(Block::GraphicControlExtension(gce));
 		}
 
 		match ib.build() {
@@ -71,9 +71,8 @@ impl GifBuilder {
 		self
 	}*/
 
-	pub fn repeat(mut self, count: u16) -> Self {
-		self.blocks
-			.push(Block::Extension(Extension::Looping(count)));
+	pub fn repeat(mut self, count: LoopCount) -> Self {
+		self.blocks.push(Block::LoopingExtension(count));
 		self
 	}