MyStitcher/src/main.rs

use fast_image_resize as fr;
use fr::{PixelType, ResizeAlg, ResizeOptions, Resizer, images::Image};
use futures::future::join_all;
use std::{
    cell::UnsafeCell,
    env,
    io::Cursor,
    ops::{Bound, RangeBounds},
    path::Path,
};
use tokio::{
    fs,
    io::{AsyncReadExt, AsyncWriteExt},
};
use tokio::{io::BufReader, net::TcpListener};

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    dotenv::dotenv().ok();
    let listener = TcpListener::bind("127.0.0.1:8080").await?;
    println!("Server listening on port 8080");
    loop {
        let (socket, _) = listener.accept().await?;
        tokio::spawn(async move {
            let timer_start = std::time::Instant::now();
            if let Err(e) = handle_client(socket).await {
                println!("Error: {}", e);
            } else {
                println!("Data sent");
            }
            let duration = timer_start.elapsed();
            println!("Operation completed in: {:.2?}", duration);
        });
    }
}

async fn handle_client(
    mut socket: tokio::net::TcpStream,
) -> Result<(), Box<dyn std::error::Error>> {
    println!("Incoming request");
    let mut reader = BufReader::new(&mut socket);

    let rect_len = reader.read_u16_le().await?;
    let mut rect_buf = vec![0u8; rect_len as usize];
    reader.read_exact(&mut rect_buf).await?;
    let rect = String::from_utf8(rect_buf)?;

    let offset_x = reader.read_f32_le().await?;
    let size_x = reader.read_f32_le().await?;
    let offset_y = reader.read_f32_le().await?;
    let size_y = reader.read_f32_le().await?;
    let scale = reader.read_f32_le().await?;

    println!("Parsed values:");
    println!("  rect:     {rect}");
    println!("  offset_x: {offset_x}");
    println!("  size_x:   {size_x}");
    println!("  offset_y: {offset_y}");
    println!("  size_y:   {size_y}");
    println!("  scale:    {scale}");

    let mut result_buffer = Vec::new();
    stitch_and_resize_image(
        &mut result_buffer,
        &rect,
        scale,
        offset_x,
        offset_y,
        size_x,
        size_y,
    )
    .await;

    socket.write_all(&result_buffer).await.unwrap();

    Ok(())
}
async fn stitch_and_resize_image(
    out: &mut Vec<u8>,
    rect: &str,
    scale: f32,
    offset_x: f32,
    offset_y: f32,
    size_x: f32,
    size_y: f32,
) {
    let mut rect = rect.split(":");
    const SIZE: u32 = 720;
    let actual_size = (SIZE as f32 * scale).trunc() as u32;

    let start = rect.next().unwrap();
    let end = rect.next().unwrap();

    let (start_col_s, start_row) = split_before_number(start);
    let start_col = column_to_number(start_col_s);
    let start_row = start_row.parse::<u32>().unwrap();
    let (end_col_s, end_row) = split_before_number(end);
    let end_col = column_to_number(end_col_s);
    let end_row = end_row.parse::<u32>().unwrap();
    let col_count = end_col - start_col + 1;
    let row_count = end_row - start_row + 1;
    let canvas_width = (row_count * actual_size) as usize;
    let canvas_height = (col_count * actual_size) as usize;
    let canvas = UnsafeVec::new(vec![0u8; canvas_width * canvas_height * 3]);

    let min_row = start_row + (row_count as f32 * offset_x).floor() as u32;
    let max_row = start_row + (row_count as f32 * (offset_x + size_x)).floor() as u32;
    let min_col = start_col + (col_count as f32 * offset_y).floor() as u32;
    let max_col = start_col + (col_count as f32 * (offset_y + size_y)).floor() as u32;

    let mut handles = Vec::new();
    for i in start_col..=end_col {
        for j in start_row..=end_row {
            if i < min_col {
                continue;
            }
            if i > max_col {
                continue;
            }
            if j < min_row {
                continue;
            }
            if j > max_row {
                continue;
            }
            handles.push(add_to_canvas(
                &canvas,
                i,
                start_col,
                j,
                start_row,
                SIZE,
                actual_size,
                canvas_width,
                canvas_height,
            ));
        }
    }
    let results = join_all(handles).await;

    let start_x = (offset_x * canvas_width as f32).trunc() as usize;
    let start_y = (offset_y * canvas_height as f32).trunc() as usize;
    let end_x = ((offset_x + size_x) * canvas_width as f32).trunc() as usize;
    let end_y = ((offset_y + size_y) * canvas_height as f32).trunc() as usize;

    let crop_width = end_x - start_x;
    let crop_height = end_y - start_y;

    let mut encoder = png::Encoder::new(
        Cursor::new(out),
        crop_width.try_into().unwrap(),
        crop_height.try_into().unwrap(),
    );
    encoder.set_color(png::ColorType::Rgb);
    encoder.set_depth(png::BitDepth::Eight);
    encoder
        .write_header()
        .unwrap()
        .write_image_data(&crop_rgb_image(
            &canvas.to_vec(),
            (actual_size * row_count) as usize,
            (actual_size * col_count) as usize,
            start_x,
            start_y,
            crop_width,
            crop_height,
        ))
        .unwrap();
}
fn split_before_number(input: &str) -> (&str, &str) {
    let index = find_first_number_start(input);
    (&input[..index], &input[index..])
}
fn find_first_number_start(s: &str) -> usize {
    for (i, c) in s.char_indices() {
        if c.is_ascii_digit() {
            return i;
        }
    }
    s.len()
}
fn number_to_column(mut num: u32) -> String {
    let mut result = String::new();

    while num > 0 {
        num -= 1;
        let ch = ((num % 26) as u8 + b'A') as char;
        result.insert(0, ch);
        num /= 26;
    }

    result
}
fn column_to_number(s: &str) -> u32 {
    let mut result = 0;

    for ch in s.chars() {
        result = result * 26 + (ch as u32 - 'A' as u32 + 1);
    }

    result
}
async fn add_to_canvas(
    canvas: &UnsafeVec<u8>,
    i: u32,
    start_col: u32,
    j: u32,
    start_row: u32,
    size: u32,
    actual_size: u32,
    canvas_width: usize,
    canvas_height: usize,
) {
    let mut resizer = Resizer::new();
    let env = &env::var("ASSET_PATH_RO").unwrap();
    let img_path = Path::new(&env);
    let filename = format!("{}{}.png", number_to_column(i), j);

    let data = Cursor::new(fs::read(img_path.join(filename)).await.unwrap());
    let decoder = png::Decoder::new(data);
    let mut reader = decoder.read_info().unwrap();
    let mut buf = vec![0; reader.output_buffer_size()];
    let _info = reader.next_frame(&mut buf).unwrap();

    let src_img = Image::from_vec_u8(size, size, buf, PixelType::U8x3).unwrap();
    let mut img = Image::new(actual_size, actual_size, PixelType::U8x3);

    resizer
        .resize(
            &src_img,
            &mut img,
            &ResizeOptions::new()
                .resize_alg(ResizeAlg::Nearest)
                .use_alpha(false),
        )
        .unwrap();
    place_image(
        canvas,
        canvas_width,
        canvas_height,
        &img.into_vec(),
        (actual_size) as usize,
        (actual_size) as usize,
        ((j - start_row) * actual_size) as usize,
        ((i - start_col) * actual_size) as usize,
    );
}
fn place_image(
    canvas: &UnsafeVec<u8>,
    canvas_width: usize,
    canvas_height: usize,
    image: &[u8],
    image_width: usize,
    image_height: usize,
    x_offset: usize,
    y_offset: usize,
) {
    let channels = 3;

    for row in 0..image_height {
        let canvas_y = y_offset + row;
        if canvas_y >= canvas_height {
            continue;
        }

        let canvas_start = (canvas_y * canvas_width + x_offset) * channels;
        let canvas_end = canvas_start + image_width * channels;

        let image_start = row * image_width * channels;
        let image_end = image_start + image_width * channels;

        if canvas_end <= canvas.len() && image_end <= image.len() {
            unsafe {
                canvas.range_mut(canvas_start..canvas_end, &image[image_start..image_end]);
            }
        }
    }
}

fn crop_rgb_image(
    img: &[u8],
    width: usize,
    height: usize,
    start_x: usize,
    start_y: usize,
    crop_width: usize,
    crop_height: usize,
) -> Vec<u8> {
    let mut cropped = Vec::with_capacity(crop_width * crop_height * 3);

    let end_x = start_x + crop_width;
    let end_y = start_y + crop_height;
    for y in start_y..end_y {
        let row_start = (y * width + start_x) * 3;
        let row_end = (y * width + end_x) * 3;
        cropped.extend_from_slice(&img[row_start..row_end]);
    }

    cropped
}

pub struct UnsafeVec<T> {
    data: UnsafeCell<Vec<T>>,
}

impl<T: Copy> UnsafeVec<T> {
    pub fn new(initial: Vec<T>) -> Self {
        UnsafeVec {
            data: UnsafeCell::new(initial),
        }
    }

    pub fn len(&self) -> usize {
        unsafe { (*self.data.get()).len() }
    }

    pub fn to_vec(self) -> Vec<T> {
        self.data.into_inner()
    }

    pub unsafe fn range_mut<R: RangeBounds<usize>>(&self, range: R, value: &[T]) {
        unsafe {
            let vec = &mut *self.data.get();

            let start = match range.start_bound() {
                Bound::Included(&s) => s,
                Bound::Excluded(&s) => s + 1,
                Bound::Unbounded => 0,
            };

            let end = match range.end_bound() {
                Bound::Included(&e) => e + 1,
                Bound::Excluded(&e) => e,
                Bound::Unbounded => vec.len(),
            };

            &mut vec[start..end].copy_from_slice(value);
        }
    }
}

unsafe impl<T> Sync for UnsafeVec<T> {}
unsafe impl<T> Send for UnsafeVec<T> {}