From 55c65fe30e7a92a4e093556baffac953c0cb9f55 Mon Sep 17 00:00:00 2001
From: ef3d0c3e <ef3d0c3e@pundalik.org>
Date: Fri, 26 Jul 2024 11:21:00 +0200
Subject: [PATCH] Partial References

---
 src/compiler/compiler.rs | 154 +++++++++++++----------
 src/document/document.rs | 167 ++++++++++++++-----------
 src/document/element.rs  | 117 +++++++++++------
 src/elements/media.rs    | 102 ++++++++++-----
 src/elements/section.rs  | 264 ++++++++++++++++++++++-----------------
 src/main.rs              |  19 ++-
 src/parser/langparser.rs |   5 +
 7 files changed, 500 insertions(+), 328 deletions(-)

diff --git a/src/compiler/compiler.rs b/src/compiler/compiler.rs
index 2ba9480..77adcbc 100644
--- a/src/compiler/compiler.rs
+++ b/src/compiler/compiler.rs
@@ -1,132 +1,155 @@
-use std::{cell::{RefCell, RefMut}, rc::Rc};
-
+use std::cell::RefCell;
+use std::cell::RefMut;
+use std::collections::HashMap;
+use std::rc::Rc;
 
 use rusqlite::Connection;
 
-use crate::document::{document::Document, variable::Variable};
+use crate::document::document::Document;
+use crate::document::document::ElemReference;
+use crate::document::element::ReferenceableElement;
+use crate::document::variable::Variable;
 
 #[derive(Clone, Copy)]
-pub enum Target
-{
+pub enum Target {
 	HTML,
 	LATEX,
 }
 
-pub struct Compiler
-{
+pub struct Compiler {
 	target: Target,
 	cache: Option<RefCell<Connection>>,
+	reference_count: RefCell<HashMap<String, HashMap<String, usize>>>,
 }
 
-impl Compiler
-{
+impl Compiler {
 	pub fn new(target: Target, db_path: Option<String>) -> Self {
-		let cache = match db_path
-		{
+		let cache = match db_path {
 			None => None,
-			Some(path) => {
-				match Connection::open(path)
-				{
-					Err(e) => panic!("Cannot connect to database: {e}"),
-					Ok(con) => Some(con),
-				}
-			}
+			Some(path) => match Connection::open(path) {
+				Err(e) => panic!("Cannot connect to database: {e}"),
+				Ok(con) => Some(con),
+			},
 		};
 		Self {
 			target,
 			cache: cache.map(|con| RefCell::new(con)),
+			reference_count: RefCell::new(HashMap::new()),
 		}
 	}
 
-	pub fn target(&self) -> Target
-	{
-		self.target
+	/// Inserts or get a reference id for the compiled document
+	///
+	/// # Parameters
+	/// - [`reference`] The reference to get or insert
+	pub fn reference_id<'a>(&self, document: &'a dyn Document, reference: ElemReference) -> usize {
+		let mut borrow = self.reference_count.borrow_mut();
+		let reference = document.get_from_reference(&reference).unwrap();
+		let refkey = reference.refcount_key();
+		let refname = reference.reference_name().unwrap();
+
+		let map = match borrow.get_mut(refkey) {
+			Some(map) => map,
+			None => {
+				borrow.insert(refkey.to_string(), HashMap::new());
+				borrow.get_mut(refkey).unwrap()
+			}
+		};
+
+		if let Some(elem) = map.get(refname) {
+			*elem
+		} else {
+			// Insert new ref
+			let index = map
+				.iter()
+				.fold(0, |max, (name, value)| std::cmp::max(max, *value));
+			map.insert(refname.clone(), index + 1);
+			index + 1
+		}
 	}
 
-	pub fn cache(&self) -> Option<RefMut<'_, Connection>>
-	{
-		self.cache
-			.as_ref()
-			.map(RefCell::borrow_mut)
+	pub fn target(&self) -> Target { self.target }
+
+	pub fn cache(&self) -> Option<RefMut<'_, Connection>> {
+		self.cache.as_ref().map(RefCell::borrow_mut)
 	}
 
 	pub fn sanitize<S: AsRef<str>>(&self, str: S) -> String {
-		match self.target
-		{
-			Target::HTML => str.as_ref()
+		match self.target {
+			Target::HTML => str
+				.as_ref()
 				.replace("&", "&amp;")
 				.replace("<", "&lt;")
 				.replace(">", "&gt;")
 				.replace("\"", "&quot;"),
-			_ => todo!("Sanitize not implemented")
+			_ => todo!("Sanitize not implemented"),
 		}
 	}
 
-	pub fn header(&self, document: &dyn Document) -> String
-	{
-		pub fn get_variable_or_error(document: &dyn Document, var_name: &'static str) -> Option<Rc<dyn Variable>>
-		{
-			document.get_variable(var_name)
+	pub fn header(&self, document: &dyn Document) -> String {
+		pub fn get_variable_or_error(
+			document: &dyn Document,
+			var_name: &'static str,
+		) -> Option<Rc<dyn Variable>> {
+			document
+				.get_variable(var_name)
 				.and_then(|var| Some(var))
 				.or_else(|| {
-					println!("Missing variable `{var_name}` in {}", document.source().name());
+					println!(
+						"Missing variable `{var_name}` in {}",
+						document.source().name()
+					);
 					None
 				})
 		}
 
 		let mut result = String::new();
-		match self.target()
-		{
+		match self.target() {
 			Target::HTML => {
 				result += "<!DOCTYPE HTML><html><head>";
 				result += "<meta charset=\"UTF-8\">";
-				if let Some(page_title) = get_variable_or_error(document, "html.page_title")
-				{
-					result += format!("<title>{}</title>", self.sanitize(page_title.to_string())).as_str();
+				if let Some(page_title) = get_variable_or_error(document, "html.page_title") {
+					result += format!("<title>{}</title>", self.sanitize(page_title.to_string()))
+						.as_str();
 				}
 
-				if let Some(css) = document.get_variable("html.css")
-				{
-					result += format!("<link rel=\"stylesheet\" href=\"{}\">", self.sanitize(css.to_string())).as_str();
+				if let Some(css) = document.get_variable("html.css") {
+					result += format!(
+						"<link rel=\"stylesheet\" href=\"{}\">",
+						self.sanitize(css.to_string())
+					)
+					.as_str();
 				}
 				result += "</head><body>";
 
 				// TODO: TOC
 				// TODO: Author, Date, Title, Div
-			},
-			Target::LATEX => {
-
-			},
+			}
+			Target::LATEX => {}
 		}
 		result
 	}
 
-	pub fn footer(&self, _document: &dyn Document) -> String
-	{
+	pub fn footer(&self, _document: &dyn Document) -> String {
 		let mut result = String::new();
-		match self.target()
-		{
+		match self.target() {
 			Target::HTML => {
 				result += "</body></html>";
-			},
-			Target::LATEX => {
-
-			},
+			}
+			Target::LATEX => {}
 		}
 		result
 	}
-	
-	pub fn compile(&self, document: &dyn Document) -> String
-	{
+
+	pub fn compile(&self, document: &dyn Document) -> String {
 		let mut out = String::new();
 		let borrow = document.content().borrow();
 
 		// Header
 		out += self.header(document).as_str();
-		
+
 		// Body
-		for i in 0 .. borrow.len()
-		{
+		for i in 0..borrow.len() {
 			let elem = &borrow[i];
 			//let prev = match i
 			//{
@@ -135,13 +158,12 @@ impl Compiler
 			//};
 			//let next = borrow.get(i+1);
 
-			match elem.compile(self, document)
-			{
+			match elem.compile(self, document) {
 				Ok(result) => {
 					//println!("Elem: {}\nCompiled to: {result}", elem.to_string());
 					out.push_str(result.as_str())
-				},
-				Err(err) => println!("Unable to compile element: {err}\n{}", elem.to_string())
+				}
+				Err(err) => println!("Unable to compile element: {err}\n{}", elem.to_string()),
 			}
 		}
 
diff --git a/src/document/document.rs b/src/document/document.rs
index 60282a1..6bbb1cd 100644
--- a/src/document/document.rs
+++ b/src/document/document.rs
@@ -7,16 +7,22 @@ use std::rc::Rc;
 use crate::parser::source::Source;
 
 use super::element::Element;
+use super::element::ReferenceableElement;
 use super::variable::Variable;
 
-// TODO: Referenceable rework
-// Usize based referencing is not an acceptable method
-// if we want to support deltas for the lsp
+#[derive(Debug, Clone, Copy)]
+pub enum ElemReference {
+	Direct(usize),
+
+	// Reference nested inside another element, e.g [`Paragraph`] or [`Media`]
+	Nested(usize, usize),
+}
+
 #[derive(Debug)]
 pub struct Scope {
 	/// List of all referenceable elements in current scope.
 	/// All elements in this should return a non empty
-	pub referenceable: HashMap<String, usize>,
+	pub referenceable: HashMap<String, ElemReference>,
 	pub variables: HashMap<String, Rc<dyn Variable>>,
 }
 
@@ -32,12 +38,16 @@ impl Scope {
 		match merge_as.is_empty() {
 			true => {
 				// References
-				self.referenceable.extend(
-					other
-						.referenceable
-						.drain()
-						.map(|(name, idx)| (name, idx + ref_offset)),
-				);
+				self.referenceable.extend(other.referenceable.drain().map(
+					|(name, idx)| match idx {
+						ElemReference::Direct(index) => {
+							(name, ElemReference::Direct(index + ref_offset))
+						}
+						ElemReference::Nested(index, sub_index) => {
+							(name, ElemReference::Nested(index + ref_offset, sub_index))
+						}
+					},
+				));
 
 				// Variables
 				self.variables
@@ -45,12 +55,18 @@ impl Scope {
 			}
 			false => {
 				// References
-				self.referenceable.extend(
-					other
-						.referenceable
-						.drain()
-						.map(|(name, idx)| (format!("{merge_as}.{name}"), idx + ref_offset)),
-				);
+				self.referenceable.extend(other.referenceable.drain().map(
+					|(name, idx)| match idx {
+						ElemReference::Direct(index) => (
+							format!("{merge_as}.{name}"),
+							ElemReference::Direct(index + ref_offset),
+						),
+						ElemReference::Nested(index, sub_index) => (
+							format!("{merge_as}.{name}"),
+							ElemReference::Nested(index + ref_offset, sub_index),
+						),
+					},
+				));
 
 				// Variables
 				self.variables.extend(
@@ -80,44 +96,41 @@ pub trait Document<'a>: core::fmt::Debug {
 
 	/// Pushes a new element into the document's content
 	fn push(&self, elem: Box<dyn Element>) {
-		// TODO: RefTable
+		if let Some(refname) = elem
+			.as_referenceable()
+			.and_then(|reference| reference.reference_name())
+		{
+			self.scope().borrow_mut().referenceable.insert(
+				refname.clone(),
+				ElemReference::Direct(self.content().borrow().len()),
+			);
+		} else if let Some(container) = self
+			.content()
+			.borrow()
+			.last()
+			.and_then(|elem| elem.as_container())
+		{
+			// This is a hack that works thanks to the fact that at document end, a [`DocumentEnd`] elem is pushed
+			container
+				.contained()
+				.iter()
+				.enumerate()
+				.for_each(|(sub_idx, elem)| {
+					if let Some(refname) = elem
+						.as_referenceable()
+						.and_then(|reference| reference.reference_name())
+					{
+						self.scope().borrow_mut().referenceable.insert(
+							refname.clone(),
+							ElemReference::Nested(self.content().borrow().len() - 1, sub_idx),
+						);
+					}
+				});
+		}
 
 		self.content().borrow_mut().push(elem);
 	}
 
-	/*
-	fn last_element(&'a self, recurse: bool) -> Option<Ref<'_, dyn Element>>
-	{
-		let elem = Ref::filter_map(self.content().borrow(),
-		|content| content.last()
-			.and_then(|last| last.downcast_ref::<Element>())
-		).ok();
-
-
-		if elem.is_some() || !recurse { return elem }
-
-		match self.parent()
-		{
-			None => None,
-			Some(parent) => parent.last_element(true),
-		}
-	}
-
-	fn last_element_mut(&'a self, recurse: bool) -> Option<RefMut<'_, dyn Element>>
-	{
-		let elem = RefMut::filter_map(self.content().borrow_mut(),
-		|content| content.last_mut()).ok();
-
-		if elem.is_some() || !recurse { return elem }
-
-		match self.parent()
-		{
-			None => None,
-			Some(parent) => parent.last_element_mut(true),
-		}
-	}
-	*/
-
 	fn add_variable(&self, variable: Rc<dyn Variable>) {
 		self.scope()
 			.borrow_mut()
@@ -141,27 +154,11 @@ pub trait Document<'a>: core::fmt::Debug {
 		}
 	}
 
-	/*
-	fn remove_variable(&self, name: &str) -> Option<Rc<dyn Variable>>
-	{
-		match self.scope().borrow_mut().variables.remove(name)
-		{
-			Some(variable) => {
-				return Some(variable.clone());
-			},
-
-			// Continue search recursively
-			None => match self.parent() {
-				Some(parent) => return parent.remove_variable(name),
-
-				// Not found
-				None => return None,
-			}
-		}
-	}
-	*/
-
 	/// Merges [`other`] into [`self`]
+	///
+	/// # Parameters
+	///
+	/// If [`merge_as`] is None, references and variables from the other document are not merged into self
 	fn merge(
 		&self,
 		content: &RefCell<Vec<Box<dyn Element>>>,
@@ -182,6 +179,36 @@ pub trait Document<'a>: core::fmt::Debug {
 			.borrow_mut()
 			.extend((content.borrow_mut()).drain(..).map(|value| value));
 	}
+
+	fn get_reference(&self, refname: &str) -> Option<ElemReference> {
+		self.scope()
+			.borrow()
+			.referenceable
+			.get(refname)
+			.and_then(|reference| Some(*reference))
+	}
+
+	fn get_from_reference(
+		&self,
+		reference: &ElemReference,
+	) -> Option<Ref<'_, dyn ReferenceableElement>> {
+		match reference {
+			ElemReference::Direct(idx) => Ref::filter_map(self.content().borrow(), |content| {
+				content.get(*idx).and_then(|elem| elem.as_referenceable())
+			})
+			.ok(),
+			ElemReference::Nested(idx, sub_idx) => {
+				Ref::filter_map(self.content().borrow(), |content| {
+					content
+						.get(*idx)
+						.and_then(|elem| elem.as_container())
+						.and_then(|container| container.contained().get(*sub_idx))
+						.and_then(|elem| elem.as_referenceable())
+				})
+				.ok()
+			}
+		}
+	}
 }
 
 pub trait DocumentAccessors<'a> {
diff --git a/src/document/element.rs b/src/document/element.rs
index c12cf98..d51c0c8 100644
--- a/src/document/element.rs
+++ b/src/document/element.rs
@@ -1,65 +1,108 @@
 use std::str::FromStr;
 
-use downcast_rs::{impl_downcast, Downcast};
-use crate::{compiler::compiler::Compiler, parser::source::Token};
+use crate::compiler::compiler::Compiler;
+use crate::parser::source::Token;
+use downcast_rs::impl_downcast;
+use downcast_rs::Downcast;
 
 use super::document::Document;
 
 #[derive(Debug, PartialEq, Eq, Clone)]
 pub enum ElemKind {
-    /// An invisible element (i.e comment)
-    Invisible,
-    /// Special elements don't trigger special formatting events
-    Special,
-    /// Inline elements don't break paragraphing
-    Inline,
-    /// Block elements are outside of paragraphs
-    Block,
+	/// An invisible element (i.e comment)
+	Invisible,
+	/// Special elements don't trigger special formatting events
+	Special,
+	/// Inline elements don't break paragraphing
+	Inline,
+	/// Block elements are outside of paragraphs
+	Block,
 }
 
 impl FromStr for ElemKind {
-    type Err = String;
+	type Err = String;
 
-    fn from_str(s: &str) -> Result<Self, Self::Err> {
-		match s
-		{
+	fn from_str(s: &str) -> Result<Self, Self::Err> {
+		match s {
 			"invisible" => Ok(ElemKind::Invisible),
-			"special"   => Ok(ElemKind::Special),
-			"inline"    => Ok(ElemKind::Inline),
-			"block"     => Ok(ElemKind::Block),
-			_ => Err(format!("Unknown ElemKind: {s}"))
+			"special" => Ok(ElemKind::Special),
+			"inline" => Ok(ElemKind::Inline),
+			"block" => Ok(ElemKind::Block),
+			_ => Err(format!("Unknown ElemKind: {s}")),
 		}
-    }
+	}
 }
 
-pub trait Element: Downcast
-{
-    /// Gets the element defined location i.e token without filename
-    fn location(&self) -> &Token;
+pub trait Element: Downcast {
+	/// Gets the element defined location i.e token without filename
+	fn location(&self) -> &Token;
 
-    fn kind(&self) -> ElemKind;
+	fn kind(&self) -> ElemKind;
 
-    /// Get the element's name
+	/// Get the element's name
 	fn element_name(&self) -> &'static str;
 
-    /// Outputs element to string for debug purposes
-    fn to_string(&self) -> String;
+	/// Outputs element to string for debug purposes
+	fn to_string(&self) -> String;
 
-    fn as_referenceable(&self) -> Option<&dyn ReferenceableElement> { None }
+	/// Gets the element as a referenceable i.e an element that can be referenced
+	fn as_referenceable(&self) -> Option<&dyn ReferenceableElement> { None }
 
-    /// Compiles element
-    fn compile(&self, compiler: &Compiler, document: &dyn Document) -> Result<String, String>;
+	/// Gets the element as a container containing other elements
+	fn as_container(&self) -> Option<&dyn ContainerElement> { None }
+
+	/// Compiles element
+	fn compile(&self, compiler: &Compiler, document: &dyn Document) -> Result<String, String>;
 }
 impl_downcast!(Element);
 
-impl core::fmt::Debug for dyn Element
-{
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        write!(f, "{}", self.to_string())
-    }
+impl core::fmt::Debug for dyn Element {
+	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+		write!(f, "{}", self.to_string())
+	}
 }
 
-pub trait ReferenceableElement : Element {
-    /// Reference name
+pub trait ReferenceableElement: Element {
+	/// Reference name
 	fn reference_name(&self) -> Option<&String>;
+
+	/// Key for refcounting
+	fn refcount_key(&self) -> &'static str;
+}
+
+impl core::fmt::Debug for dyn ReferenceableElement {
+	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+		write!(f, "{}", self.to_string())
+	}
+}
+
+pub trait ContainerElement: Element {
+	/// Gets the contained elements
+	fn contained(&self) -> &Vec<Box<dyn Element>>;
+
+	/// Adds an element to the container
+	fn push(&mut self, elem: Box<dyn Element>) -> Result<(), String>;
+}
+
+impl core::fmt::Debug for dyn ContainerElement {
+	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+		write!(f, "{}", self.to_string())
+	}
+}
+
+#[derive(Debug)]
+pub struct DocumentEnd(pub Token);
+
+impl Element for DocumentEnd {
+	fn location(&self) -> &Token { &self.0 }
+
+	fn kind(&self) -> ElemKind { ElemKind::Invisible }
+
+	fn element_name(&self) -> &'static str { "Document End" }
+
+	fn to_string(&self) -> String { format!("{self:#?}") }
+
+	fn compile(&self, _compiler: &Compiler, _document: &dyn Document) -> Result<String, String> {
+		Ok(String::new())
+	}
 }
diff --git a/src/elements/media.rs b/src/elements/media.rs
index 0dfe144..1923cb4 100644
--- a/src/elements/media.rs
+++ b/src/elements/media.rs
@@ -16,8 +16,10 @@ use crate::compiler::compiler::Compiler;
 use crate::compiler::compiler::Target;
 use crate::document::document::Document;
 use crate::document::document::DocumentAccessors;
+use crate::document::element::ContainerElement;
 use crate::document::element::ElemKind;
 use crate::document::element::Element;
+use crate::document::element::ReferenceableElement;
 use crate::document::references::validate_refname;
 use crate::parser::parser::ReportColors;
 use crate::parser::rule::RegexRule;
@@ -54,36 +56,22 @@ impl FromStr for MediaType {
 }
 
 #[derive(Debug)]
-struct MediaGroup {
+struct Media {
 	pub(self) location: Token,
-	pub(self) media: Vec<Media>,
+	pub(self) media: Vec<Box<dyn Element>>,
 }
 
-impl MediaGroup {
-	fn push(&mut self, media: Media) -> Result<(), String> {
-		if self.location.source() != media.location.source() {
-			return Err(format!(
-				"Attempted to insert media from {} into MediaGroup from {}",
-				self.location.source(),
-				media.location.source()
-			));
-		}
-
-		self.location.range = self.location.start()..media.location.end();
-		self.media.push(media);
-		Ok(())
-	}
-}
-
-impl Element for MediaGroup {
+impl Element for Media {
 	fn location(&self) -> &Token { &self.location }
 
 	fn kind(&self) -> ElemKind { ElemKind::Block }
 
-	fn element_name(&self) -> &'static str { "Media Group" }
+	fn element_name(&self) -> &'static str { "Media" }
 
 	fn to_string(&self) -> String { format!("{self:#?}") }
 
+	fn as_container(&self) -> Option<&dyn ContainerElement> { Some(self) }
+
 	fn compile(&self, compiler: &Compiler, document: &dyn Document) -> Result<String, String> {
 		match compiler.target() {
 			Target::HTML => {
@@ -105,8 +93,30 @@ impl Element for MediaGroup {
 	}
 }
 
+impl ContainerElement for Media {
+	fn contained(&self) -> &Vec<Box<dyn Element>> { &self.media }
+
+	fn push(&mut self, elem: Box<dyn Element>) -> Result<(), String> {
+		let medium = match elem.downcast_ref::<Medium>() {
+			Some(medium) => medium,
+			None => return Err("Attempted to insert invalid element into Media".to_string()),
+		};
+		if self.location.source() != medium.location.source() {
+			return Err(format!(
+				"Attempted to insert medium from {} into medium from {}",
+				self.location.source(),
+				medium.location.source()
+			));
+		}
+
+		self.location.range = self.location.start()..medium.location.end();
+		self.media.push(elem);
+		Ok(())
+	}
+}
+
 #[derive(Debug)]
-struct Media {
+struct Medium {
 	pub(self) location: Token,
 	pub(self) reference: String,
 	pub(self) uri: String,
@@ -116,15 +126,17 @@ struct Media {
 	pub(self) description: Option<Paragraph>,
 }
 
-impl Element for Media {
+impl Element for Medium {
 	fn location(&self) -> &Token { &self.location }
 
 	fn kind(&self) -> ElemKind { ElemKind::Block }
 
-	fn element_name(&self) -> &'static str { "Media" }
+	fn element_name(&self) -> &'static str { "Medium" }
 
 	fn to_string(&self) -> String { format!("{self:#?}") }
 
+	fn as_referenceable(&self) -> Option<&dyn ReferenceableElement> { Some(self) }
+
 	fn compile(&self, compiler: &Compiler, document: &dyn Document) -> Result<String, String> {
 		match compiler.target() {
 			Target::HTML => {
@@ -142,7 +154,19 @@ impl Element for Media {
 					MediaType::VIDEO => todo!(),
 					MediaType::AUDIO => todo!(),
 				}
-				result.push_str(format!(r#"<p class="medium-refname">{}</p>"#, "TODO").as_str());
+
+				let caption = self
+					.caption
+					.as_ref()
+					.and_then(|cap| Some(format!(" {}", compiler.sanitize(cap.as_str()))))
+					.unwrap_or(String::new());
+
+				// Reference
+				let elemref = document.get_reference(self.reference.as_str()).unwrap();
+				let refcount = compiler.reference_id(document, elemref);
+				result.push_str(
+					format!(r#"<p class="medium-refname">({refcount}){caption}</p>"#).as_str(),
+				);
 				if let Some(paragraph) = self.description.as_ref() {
 					match paragraph.compile(compiler, document) {
 						Ok(res) => result.push_str(res.as_str()),
@@ -158,6 +182,12 @@ impl Element for Media {
 	}
 }
 
+impl ReferenceableElement for Medium {
+	fn reference_name(&self) -> Option<&String> { Some(&self.reference) }
+
+	fn refcount_key(&self) -> &'static str { "medium" }
+}
+
 pub struct MediaRule {
 	re: [Regex; 1],
 	properties: PropertyParser,
@@ -178,6 +208,10 @@ impl MediaRule {
 			"width".to_string(),
 			Property::new(false, "Override for the media width".to_string(), None),
 		);
+		props.insert(
+			"caption".to_string(),
+			Property::new(false, "Medium caption".to_string(), None),
+		);
 		Self {
 			re: [RegexBuilder::new(
 				r"^!\[(.*)\]\(((?:\\.|[^\\\\])*?)\)(?:\[((?:\\.|[^\\\\])*?)\])?((?:\\(?:.|\n)|[^\\\\])*?$)?",
@@ -366,6 +400,13 @@ impl RegexRule for MediaRule {
 			.ok()
 			.and_then(|(_, s)| Some(s));
 
+		let caption = properties
+			.get("caption", |_, value| -> Result<String, ()> {
+				Ok(value.clone())
+			})
+			.ok()
+			.and_then(|(_, value)| Some(value));
+
 		let description = match matches.get(4) {
 			Some(content) => {
 				let source = Rc::new(VirtualSource::new(
@@ -399,31 +440,30 @@ impl RegexRule for MediaRule {
 			None => panic!("Unknown error"),
 		};
 
-		// TODO: caption
-		let mut group = match document.last_element_mut::<MediaGroup>() {
+		let mut group = match document.last_element_mut::<Media>() {
 			Some(group) => group,
 			None => {
 				parser.push(
 					document,
-					Box::new(MediaGroup {
+					Box::new(Media {
 						location: token.clone(),
 						media: vec![],
 					}),
 				);
 
-				document.last_element_mut::<MediaGroup>().unwrap()
+				document.last_element_mut::<Media>().unwrap()
 			}
 		};
 
-		if let Err(err) = group.push(Media {
+		if let Err(err) = group.push(Box::new(Medium {
 			location: token.clone(),
 			reference: refname,
 			uri,
 			media_type,
 			width,
-			caption: None,
+			caption,
 			description,
-		}) {
+		})) {
 			reports.push(
 				Report::build(ReportKind::Error, token.source(), token.start())
 					.with_message("Invalid Media")
diff --git a/src/elements/section.rs b/src/elements/section.rs
index 9af30a6..3b9dbe3 100644
--- a/src/elements/section.rs
+++ b/src/elements/section.rs
@@ -1,42 +1,57 @@
-use mlua::{Error::BadArgument, Function, Lua};
+use crate::compiler::compiler::Compiler;
+use crate::compiler::compiler::Target;
+use crate::document::document::Document;
+use crate::document::element::ElemKind;
+use crate::document::element::Element;
+use crate::document::element::ReferenceableElement;
+use crate::lua::kernel::CTX;
+use crate::parser::parser::Parser;
+use crate::parser::rule::RegexRule;
+use crate::parser::source::Source;
+use crate::parser::source::Token;
+use ariadne::Fmt;
+use ariadne::Label;
+use ariadne::Report;
+use ariadne::ReportKind;
+use mlua::Error::BadArgument;
+use mlua::Function;
+use mlua::Lua;
 use regex::Regex;
-use crate::{compiler::compiler::Target, document::document::Document, lua::kernel::CTX, parser::{parser::Parser, rule::RegexRule, source::{Source, Token}}};
-use ariadne::{Report, Fmt, Label, ReportKind};
-use crate::{compiler::compiler::Compiler, document::element::{ElemKind, Element, ReferenceableElement}};
-use std::{ops::Range, rc::Rc, sync::Arc};
+use std::ops::Range;
+use std::rc::Rc;
+use std::sync::Arc;
 
 #[derive(Debug)]
 pub struct Section {
-    pub(self) location: Token,
-	pub(self) title: String, // Section title
-	pub(self) depth: usize, // Section depth
-	pub(self) kind: u8, // Section kind, e.g numbered, unnumbred, ...
+	pub(self) location: Token,
+	pub(self) title: String,             // Section title
+	pub(self) depth: usize,              // Section depth
+	pub(self) kind: u8,                  // Section kind, e.g numbered, unnumbred, ...
 	pub(self) reference: Option<String>, // Section reference name
 }
 
-impl Element for Section
-{
-    fn location(&self) -> &Token { &self.location }
-    fn kind(&self) -> ElemKind { ElemKind::Block }
-    fn element_name(&self) -> &'static str { "Section" }
-    fn to_string(&self) -> String { format!("{self:#?}") }
-    fn as_referenceable(&self) -> Option<&dyn ReferenceableElement> { Some(self) }
-    fn compile(&self, compiler: &Compiler, _document: &dyn Document) -> Result<String, String> {
-        match compiler.target()
-        {
-            Target::HTML => {
-                Ok(format!("<h{0}>{1}</h{0}>",
-                        self.depth,
-                        compiler.sanitize(self.title.as_str())))
-            }
-            Target::LATEX => Err("Unimplemented compiler".to_string())
-        }
-    }
+impl Element for Section {
+	fn location(&self) -> &Token { &self.location }
+	fn kind(&self) -> ElemKind { ElemKind::Block }
+	fn element_name(&self) -> &'static str { "Section" }
+	fn to_string(&self) -> String { format!("{self:#?}") }
+	fn as_referenceable(&self) -> Option<&dyn ReferenceableElement> { Some(self) }
+	fn compile(&self, compiler: &Compiler, _document: &dyn Document) -> Result<String, String> {
+		match compiler.target() {
+			Target::HTML => Ok(format!(
+				"<h{0}>{1}</h{0}>",
+				self.depth,
+				compiler.sanitize(self.title.as_str())
+			)),
+			Target::LATEX => Err("Unimplemented compiler".to_string()),
+		}
+	}
 }
 
-impl ReferenceableElement for Section
-{
-    fn reference_name(&self) -> Option<&String> { self.reference.as_ref() }
+impl ReferenceableElement for Section {
+	fn reference_name(&self) -> Option<&String> { self.reference.as_ref() }
+
+	fn refcount_key(&self) -> &'static str { "section" }
 }
 
 pub struct SectionRule {
@@ -45,15 +60,16 @@ pub struct SectionRule {
 
 impl SectionRule {
 	pub fn new() -> Self {
-		Self { re: [Regex::new(r"(?:^|\n)(#{1,})(?:\{(.*)\})?((\*|\+){1,})?(.*)").unwrap()] }
+		Self {
+			re: [Regex::new(r"(?:^|\n)(#{1,})(?:\{(.*)\})?((\*|\+){1,})?(.*)").unwrap()],
+		}
 	}
 }
 
-pub mod section_kind
-{
-	pub const NONE : u8 = 0x00;
-	pub const NO_TOC : u8 = 0x01;
-	pub const NO_NUMBER : u8 = 0x02;
+pub mod section_kind {
+	pub const NONE: u8 = 0x00;
+	pub const NO_TOC: u8 = 0x01;
+	pub const NO_NUMBER: u8 = 0x02;
 }
 
 impl RegexRule for SectionRule {
@@ -61,11 +77,16 @@ impl RegexRule for SectionRule {
 
 	fn regexes(&self) -> &[Regex] { &self.re }
 
-	fn on_regex_match(&self, _: usize, parser: &dyn Parser, document: &dyn Document, token: Token, matches: regex::Captures) -> Vec<Report<'_, (Rc<dyn Source>, Range<usize>)>>
-	{
+	fn on_regex_match(
+		&self,
+		_: usize,
+		parser: &dyn Parser,
+		document: &dyn Document,
+		token: Token,
+		matches: regex::Captures,
+	) -> Vec<Report<'_, (Rc<dyn Source>, Range<usize>)>> {
 		let mut result = vec![];
-		let section_depth = match matches.get(1)
-		{
+		let section_depth = match matches.get(1) {
 			Some(depth) => {
 				if depth.len() > 6 {
 					result.push(
@@ -78,16 +99,17 @@ impl RegexRule for SectionRule {
                             6.fg(parser.colors().info)))
 							.with_color(parser.colors().error))
 						.finish());
-						return result;
+					return result;
 				}
 
 				depth.len()
-			},
+			}
 			_ => panic!("Empty section depth"),
 		};
 
 		// [Optional] Reference name
-		let section_refname = matches.get(2).map_or_else(|| None,
+		let section_refname = matches.get(2).map_or_else(
+			|| None,
 			|refname| {
 				/* TODO: Wait for reference rework
 				// Check for duplicate reference
@@ -105,29 +127,28 @@ impl RegexRule for SectionRule {
 									refname.as_str().fg(parser.colors().highlight),
 									reference.element_name().fg(parser.colors().highlight)))
 							.with_color(parser.colors().warning))
-                        .with_label(
-                            Label::new((ref_doc.source(), reference.location().start()+1..reference.location().end() ))
-                            .with_message(format!("`{}` previously defined here",
-                                refname.as_str().fg(parser.colors().highlight)))
+						.with_label(
+							Label::new((ref_doc.source(), reference.location().start()+1..reference.location().end() ))
+							.with_message(format!("`{}` previously defined here",
+								refname.as_str().fg(parser.colors().highlight)))
 							.with_color(parser.colors().warning))
 						.with_note(format!("Previous reference was overwritten"))
 						.finish());
 				}
 				*/
-			Some(refname.as_str().to_string())
-		});
+				Some(refname.as_str().to_string())
+			},
+		);
 
 		// Section kind
-		let section_kind = match matches.get(3)
-		{
-			Some(kind) => {
-				match kind.as_str() {
-					"*+" | "+*" => section_kind::NO_NUMBER | section_kind::NO_TOC,
-					"*" => section_kind::NO_NUMBER,
-					"+" => section_kind::NO_TOC,
-					"" => section_kind::NONE,
-					_ => {
-						result.push(
+		let section_kind = match matches.get(3) {
+			Some(kind) => match kind.as_str() {
+				"*+" | "+*" => section_kind::NO_NUMBER | section_kind::NO_TOC,
+				"*" => section_kind::NO_NUMBER,
+				"+" => section_kind::NO_TOC,
+				"" => section_kind::NONE,
+				_ => {
+					result.push(
 							Report::build(ReportKind::Error, token.source(), kind.start())
 							.with_message("Invalid section numbering kind")
 							.with_label(
@@ -139,38 +160,40 @@ impl RegexRule for SectionRule {
 								.with_color(parser.colors().error))
 								.with_help(format!("Leave empty for a numbered listed section"))
 							.finish());
-						return result;
-					}
+					return result;
 				}
-			}
+			},
 			_ => section_kind::NONE,
 		};
 
 		// Spacing + Section name
-		let section_name = match matches.get(5)
-		{
+		let section_name = match matches.get(5) {
 			Some(name) => {
-				let split = name.as_str().chars()
+				let split = name
+					.as_str()
+					.chars()
 					.position(|c| !c.is_whitespace())
 					.unwrap_or(0);
 
 				let section_name = &name.as_str()[split..];
-				if section_name.is_empty() // No name
+				if section_name.is_empty()
+				// No name
 				{
 					result.push(
 						Report::build(ReportKind::Error, token.source(), name.start())
-						.with_message("Missing section name")
-						.with_label(
-							Label::new((token.source(), name.range()))
-							.with_message("Sections require a name before line end")
-							.with_color(parser.colors().error))
-						.finish());
+							.with_message("Missing section name")
+							.with_label(
+								Label::new((token.source(), name.range()))
+									.with_message("Sections require a name before line end")
+									.with_color(parser.colors().error),
+							)
+							.finish(),
+					);
 					return result;
 				}
 
 				// No spacing
-				if split == 0
-				{
+				if split == 0 {
 					result.push(
 						Report::build(ReportKind::Warning, token.source(), name.start())
 						.with_message("Missing section spacing")
@@ -183,57 +206,70 @@ impl RegexRule for SectionRule {
 					return result;
 				}
 
-
 				section_name.to_string()
-			},
-			_ => panic!("Empty section name")
+			}
+			_ => panic!("Empty section name"),
 		};
 
-        parser.push(document, Box::new(
-            Section {
+		parser.push(
+			document,
+			Box::new(Section {
 				location: token.clone(),
-                title: section_name,
-                depth: section_depth,
-                kind: section_kind,
-                reference: section_refname
-            }
-        ));
+				title: section_name,
+				depth: section_depth,
+				kind: section_kind,
+				reference: section_refname,
+			}),
+		);
 
-        return result;
+		return result;
 	}
 
-	fn lua_bindings<'lua>(&self, lua: &'lua Lua) -> Vec<(String, Function<'lua>)>
-	{
+	fn lua_bindings<'lua>(&self, lua: &'lua Lua) -> Vec<(String, Function<'lua>)> {
 		let mut bindings = vec![];
 
-		bindings.push(("push".to_string(), lua.create_function(
-			|_, (title, depth, kind, reference) : (String, usize, String, Option<String>)| {
-			let kind = match kind.as_str() {
-				"*+" | "+*" => section_kind::NO_NUMBER | section_kind::NO_TOC,
-				"*" => section_kind::NO_NUMBER,
-				"+" => section_kind::NO_TOC,
-				"" => section_kind::NONE,
-				_ => return Err(BadArgument {
-						to: Some("push".to_string()),
-						pos: 3,
-						name: Some("kind".to_string()),
-						cause: Arc::new(mlua::Error::external(
-									format!("Unknown section kind specified")))})
-			};
+		bindings.push((
+			"push".to_string(),
+			lua.create_function(
+				|_, (title, depth, kind, reference): (String, usize, String, Option<String>)| {
+					let kind = match kind.as_str() {
+						"*+" | "+*" => section_kind::NO_NUMBER | section_kind::NO_TOC,
+						"*" => section_kind::NO_NUMBER,
+						"+" => section_kind::NO_TOC,
+						"" => section_kind::NONE,
+						_ => {
+							return Err(BadArgument {
+								to: Some("push".to_string()),
+								pos: 3,
+								name: Some("kind".to_string()),
+								cause: Arc::new(mlua::Error::external(format!(
+									"Unknown section kind specified"
+								))),
+							})
+						}
+					};
 
-			CTX.with_borrow(|ctx| ctx.as_ref().map(|ctx| {
-				ctx.parser.push(ctx.document, Box::new(Section {
-					location: ctx.location.clone(),
-					title,
-					depth,
-					kind,
-					reference
-				}));
-			}));
+					CTX.with_borrow(|ctx| {
+						ctx.as_ref().map(|ctx| {
+							ctx.parser.push(
+								ctx.document,
+								Box::new(Section {
+									location: ctx.location.clone(),
+									title,
+									depth,
+									kind,
+									reference,
+								}),
+							);
+						})
+					});
+
+					Ok(())
+				},
+			)
+			.unwrap(),
+		));
 
-			Ok(())
-		}).unwrap()));
-		
 		bindings
 	}
 }
diff --git a/src/main.rs b/src/main.rs
index b24efdd..77295e3 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -85,16 +85,15 @@ fn main() {
 		println!("-- END AST DEBUGGING --");
 	}
 
-	// TODO
-	//if debug_opts.contains(&"ref".to_string())
-	//{
-	//	println!("-- BEGIN REFERENCES DEBUGGING --");
-	//	let sc = doc.scope.borrow();
-	//	sc.referenceable.iter().for_each(|(name, pos)| {
-	//		println!(" - {name}: `{:#?}`", doc.content.borrow()[*pos]);
-	//	});
-	//	println!("-- END REFERENCES DEBUGGING --");
-	//}
+	if debug_opts.contains(&"ref".to_string())
+	{
+		println!("-- BEGIN REFERENCES DEBUGGING --");
+		let sc = doc.scope().borrow();
+		sc.referenceable.iter().for_each(|(name, reference)| {
+			println!(" - {name}: `{:#?}`", doc.get_from_reference(reference));
+		});
+		println!("-- END REFERENCES DEBUGGING --");
+	}
 	if debug_opts.contains(&"var".to_string()) {
 		println!("-- BEGIN VARIABLES DEBUGGING --");
 		let sc = doc.scope().borrow();
diff --git a/src/parser/langparser.rs b/src/parser/langparser.rs
index a4bed22..dddfac7 100644
--- a/src/parser/langparser.rs
+++ b/src/parser/langparser.rs
@@ -10,6 +10,7 @@ use ariadne::Report;
 
 use crate::document::document::Document;
 use crate::document::document::DocumentAccessors;
+use crate::document::element::DocumentEnd;
 use crate::document::element::ElemKind;
 use crate::document::element::Element;
 use crate::document::langdocument::LangDocument;
@@ -227,6 +228,10 @@ impl Parser for LangParser {
 				.on_scope_end(self, &doc, super::state::Scope::DOCUMENT),
 		);
 
+		self.push(&doc, Box::new(DocumentEnd(
+			Token::new(doc.source().content().len()..doc.source().content().len(), doc.source())
+		)));
+
 		return Box::new(doc);
 	}