From 6ee37f3ce9f85a3c48145a0c3ad5c2e3af4b1018 Mon Sep 17 00:00:00 2001
From: Matt Fellenz <matt@felle.nz>
Date: Tue, 26 Mar 2024 21:30:34 -0700
Subject: [PATCH] Allow multiple pages

---
 bot/src/bot.rs       | 20 ++++++++++----
 protocol/src/lib.rs  |  6 ++--
 worker/src/render.rs | 65 +++++++++++++++++++++++++++-----------------
 3 files changed, 56 insertions(+), 35 deletions(-)

diff --git a/bot/src/bot.rs b/bot/src/bot.rs
index b7afcb1..1eea2ac 100644
--- a/bot/src/bot.rs
+++ b/bot/src/bot.rs
@@ -324,14 +324,17 @@ async fn render(
 
 	match res {
 		Ok(res) => {
-			let image = CreateAttachment::bytes(res.image, "rendered.png");
-			let mut message = CreateReply::default().attachment(image).reply(true);
+			let mut message = CreateReply::default().reply(true);
 
 			let mut content = String::new();
 
-			if let Some(more_pages) = res.more_pages {
-				let more_pages = more_pages.get();
-				write!(
+			if res.images.is_empty() {
+				writeln!(content, "Note: no pages generated").unwrap();
+			}
+
+			if res.more_pages > 0 {
+				let more_pages = res.more_pages;
+				writeln!(
 					content,
 					"Note: {more_pages} more page{s} ignored",
 					s = if more_pages == 1 { "" } else { "s" },
@@ -340,7 +343,7 @@ async fn render(
 			}
 
 			if !res.warnings.is_empty() {
-				write!(
+				writeln!(
 					content,
 					"Render succeeded with warnings:\n```ansi\n{}\n```",
 					sanitize_code_block(&res.warnings),
@@ -352,6 +355,11 @@ async fn render(
 				message = message.content(content);
 			}
 
+			for (i, image) in res.images.into_iter().enumerate() {
+				let image = CreateAttachment::bytes(image, format!("page-{}.png", i + 1));
+				message = message.attachment(image);
+			}
+
 			ctx.send(message).await?;
 		}
 		Err(error) => {
diff --git a/protocol/src/lib.rs b/protocol/src/lib.rs
index e9df0db..1ebe0a8 100644
--- a/protocol/src/lib.rs
+++ b/protocol/src/lib.rs
@@ -1,5 +1,3 @@
-use std::num::NonZeroUsize;
-
 use serde::{Deserialize, Serialize};
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
@@ -11,8 +9,8 @@ pub enum Request {
 
 #[derive(Debug, Serialize, Deserialize)]
 pub struct Rendered {
-	pub image: Vec<u8>,
-	pub more_pages: Option<NonZeroUsize>,
+	pub images: Vec<Vec<u8>>,
+	pub more_pages: usize,
 	pub warnings: String,
 }
 
diff --git a/worker/src/render.rs b/worker/src/render.rs
index 527b06d..d09a784 100644
--- a/worker/src/render.rs
+++ b/worker/src/render.rs
@@ -1,5 +1,4 @@
 use std::io::Cursor;
-use std::num::NonZeroUsize;
 
 use protocol::Rendered;
 use typst::eval::Tracer;
@@ -50,6 +49,9 @@ fn to_string(v: impl ToString) -> String {
 	v.to_string()
 }
 
+const PAGE_LIMIT: usize = 5;
+const BYTES_LIMIT: usize = 25 * 1024 * 1024;
+
 pub fn render(sandbox: &Sandbox, source: String) -> Result<Rendered, String> {
 	let world = sandbox.with_source(source);
 
@@ -58,34 +60,47 @@ pub fn render(sandbox: &Sandbox, source: String) -> Result<Rendered, String> {
 		typst::compile(&world, &mut tracer).map_err(|diags| format_diagnostics(&world, &diags))?;
 	let warnings = tracer.warnings();
 
-	let frame = &document
+	let transparent = Color::from_u8(0, 0, 0, 0);
+	let mut total_attachment_size = 0;
+
+	let images = document
 		.pages
-		.first()
-		.ok_or("no pages in rendered output")?
-		.frame;
-	let more_pages = NonZeroUsize::new(document.pages.len().saturating_sub(1));
+		.iter()
+		.take(PAGE_LIMIT)
+		.map(|page| {
+			let frame = &page.frame;
+			let pixels_per_point = determine_pixels_per_point(frame.size()).map_err(to_string)?;
+			let pixmap = typst_render::render(frame, pixels_per_point, transparent);
+
+			let mut writer = Cursor::new(Vec::new());
+
+			// The unwrap will never fail since `Vec`'s `Write` implementation is infallible.
+			image::write_buffer_with_format(
+				&mut writer,
+				bytemuck::cast_slice(pixmap.pixels()),
+				pixmap.width(),
+				pixmap.height(),
+				image::ColorType::Rgba8,
+				image::ImageFormat::Png,
+			)
+			.unwrap();
+
+			Ok(writer.into_inner())
+		})
+		.take_while(|image| {
+			if let Ok(image) = image {
+				total_attachment_size += image.len();
+				total_attachment_size <= BYTES_LIMIT
+			} else {
+				true
+			}
+		})
+		.collect::<Result<Vec<_>, String>>()?;
 
-	let pixels_per_point = determine_pixels_per_point(frame.size()).map_err(to_string)?;
+	let more_pages = document.pages.len() - images.len();
 
-	let transparent = Color::from_u8(0, 0, 0, 0);
-	let pixmap = typst_render::render(frame, pixels_per_point, transparent);
-
-	let mut writer = Cursor::new(Vec::new());
-
-	// The unwrap will never fail since `Vec`'s `Write` implementation is infallible.
-	image::write_buffer_with_format(
-		&mut writer,
-		bytemuck::cast_slice(pixmap.pixels()),
-		pixmap.width(),
-		pixmap.height(),
-		image::ColorType::Rgba8,
-		image::ImageFormat::Png,
-	)
-	.unwrap();
-
-	let image = writer.into_inner();
 	Ok(Rendered {
-		image,
+		images,
 		more_pages,
 		warnings: format_diagnostics(&world, &warnings),
 	})