blockの巣

Rustでコンパイル時レイトレーシング2

2022/01/31 12:35 公開
2022/01/21 20:48

Vec3AddSubが実装されていなかったのを修正

2022/01/21 21:43

リポジトリへのリンクを追加

2022/02/16 01:26

リポジトリへのリンクを修正

Compile Time Ray Tracing in Rust Ray Tracing Rust

前回記事
リポジトリ
記事の内容はコミットID82481b974f2684a73a4940a85f92de74f9f05b7eのものです。

今回はVec3の定義とピクセルのバッファの型にタプルを使っていた部分をColor(Vec3)に置き換えるところまで行います。 なのでまだレイトレーシングらしい処理は出てこず、レンダリング結果も前回と変わりません。

今回追加したファイルはlib.rs, util.rs, vec3.rsです。
lib.rsは各モジュールの宣言、useを行っており、util.rsVec3で使用する関数を定義しています。
vec3.rsVec3型とそれに関連する関数の定義を行っています。

lib.rs

まずはlib.rsの内容からです。

必要なfeatureを有効化、各モジュールの宣言、use宣言、Vec3のエイリアスとしてPoint3Colorの定義を行っています。
有効化したfeatureはTraitの関数をconst fnとして実装できるようになる(かもしれない)話const fnで使用可能なVec3を定義するで解説しています。

#![feature(const_trait_impl)]
#![feature(const_fn_floating_point_arithmetic)]
#![feature(const_mut_refs)]
pub mod util;
pub mod vec3;

pub use vec3::*;

pub type Point3 = Vec3;
pub type Color = Vec3;

util.rs

util.rsではVec3の処理で使いたい関数の定義を行っています。具体的にはsqrtabsです。
どちらも組み込みのf64の関数として定義されていますがconst fnでは使えないので自前で定義しました。

テストを除いたutil.rsのコードです。

pub const fn sqrt(s: f64) -> f64 {
    if s == 0.0 {
        return 0.0;
    }

    let mut x = s / 2.0;
    let mut _last_x = 0.0;

    while x != _last_x {
        _last_x = x;
        x = (x + s / x) / 2.0;
    }
    x
}

pub const fn abs(s: f64) -> f64 {
    if s >= 0.0 {
        s
    } else {
        -s
    }
}

2.0.sqrt()のような形式ではなくsqrt(2.0)のような形式で呼び出すように実装しています。
sqrtは標準のf64::sqrtと若干誤差が出てしまいますが実際に使う分には問題ない程度だと思うのでこのまま進めます。

Vec3.rs

vec3.rsではVec3とそれに関連する関数を定義しています。
Vec3の定義は以下のようになっており、x, y, zには外部からもアクセスできるシンプルな作りになっています。
ここにAdd, AddAssign, Mul,Indexなどの演算子と、length_squred, length dot, corss, unit_vectorなどの関数を実装しています。
(好みで言えばunit_vectorという外部の関数よりはv.normalized()のような使い方ができる方が好きなのですが、参考元がこうなっているのでこのように実装しました)

#[derive(Debug, PartialEq, Clone, Copy)]
pub struct Vec3 {
    pub x: f64,
    pub y: f64,
    pub z: f64,
}

少々長いですがテストを除いたコード全文も貼っておきます

Vec3のテストを除いたコード全文
use crate::util::*;
use std::ops::*;

#[derive(Debug, PartialEq, Copy)]
pub struct Vec3 {
    pub x: f64,
    pub y: f64,
    pub z: f64,
}

impl const std::default::Default for Vec3 {
    fn default() -> Self {
        Vec3::ZERO
    }
}

impl const Clone for Vec3 {
    fn clone(&self) -> Self {
        Vec3 {
            x: self.x,
            y: self.y,
            z: self.z,
        }
    }

    fn clone_from(&mut self, source: &Self) {
        *self = Vec3 {
            x: source.x,
            y: source.y,
            z: source.z,
        }
    }
}

impl Vec3 {
    pub const ZERO: Vec3 = Vec3::new(0.0, 0.0, 0.0);

    pub const fn new(x: f64, y: f64, z: f64) -> Self {
        Self { x, y, z }
    }

    pub const fn length_squared(&self) -> f64 {
        self.x * self.x + self.y * self.y + self.z * self.z
    }

    pub const fn length(&self) -> f64 {
        sqrt(self.length_squared())
    }
}

impl const Neg for Vec3 {
    type Output = Self;
    fn neg(self) -> Self::Output {
        Vec3::new(-self.x, -self.y, -self.z)
    }
}

impl const Index<usize> for Vec3 {
    type Output = f64;

    fn index(&self, index: usize) -> &Self::Output {
        match index {
            0 => &self.x,
            1 => &self.y,
            2 => &self.z,
            _ => panic!("out of range"),
        }
    }
}

impl const IndexMut<usize> for Vec3 {
    fn index_mut(&mut self, index: usize) -> &mut Self::Output {
        match index {
            0 => &mut self.x,
            1 => &mut self.y,
            2 => &mut self.z,
            _ => panic!("out of range"),
        }
    }
}

impl const Add<Vec3> for Vec3 {
    type Output = Self;

    fn add(self, rhs: Vec3) -> Self::Output {
        Vec3 {
            x: self.x + rhs.x,
            y: self.y + rhs.y,
            z: self.z + rhs.z,
        }
    }
}

impl const Sub<Vec3> for Vec3 {
    type Output = Self;

    fn sub(self, rhs: Vec3) -> Self::Output {
        Vec3 {
            x: self.x - rhs.x,
            y: self.y - rhs.y,
            z: self.z - rhs.z,
        }
    }
}

impl const Mul<f64> for Vec3 {
    type Output = Vec3;

    fn mul(self, rhs: f64) -> Self::Output {
        Vec3 {
            x: self.x * rhs,
            y: self.y * rhs,
            z: self.z * rhs,
        }
    }
}

impl const MulAssign<f64> for Vec3 {
    fn mul_assign(&mut self, rhs: f64) {
        *self = *self * rhs;
    }
}

impl const Div<f64> for Vec3 {
    type Output = Vec3;

    fn div(self, rhs: f64) -> Self::Output {
        Vec3 {
            x: self.x / rhs,
            y: self.y / rhs,
            z: self.z / rhs,
        }
    }
}

impl const DivAssign<f64> for Vec3 {
    fn div_assign(&mut self, rhs: f64) {
        *self = *self / rhs;
    }
}

impl const Mul<Vec3> for f64 {
    type Output = Vec3;

    fn mul(self, rhs: Vec3) -> Self::Output {
        rhs * self
    }
}

impl const Div<Vec3> for f64 {
    type Output = Vec3;

    fn div(self, rhs: Vec3) -> Self::Output {
        rhs / self
    }
}

pub const fn dot(u: &Vec3, v: &Vec3) -> f64 {
    u.x * v.x + u.y * v.y + u.z * v.z
}

pub const fn cross(u: &Vec3, v: &Vec3) -> Vec3 {
    Vec3::new(
        u.y * v.z - u.z * v.y,
        u.z * v.x - u.x * v.z,
        u.x * v.y - u.y * v.x,
    )
}

pub const fn unit_vector(v: &Vec3) -> Vec3 {
    let len = v.length();
    if std::f64::EPSILON < abs(len) {
        *v / len
    } else {
        Vec3::ZERO
    }
}

main.rsの変更点

main.rsでピクセルの色として使用していた(f64, f64, f64)Color(Vec3)型に置き換えました。
また(f64, f64, f64)から(u8, u8, u8)に変換する処理をpixels_to_buffer()関数に移動するなどの細かい変更もしています。

main.rs全文
#![feature(const_fn_floating_point_arithmetic)] // const fnの中で浮動小数点演算を行うために必要

use image::ImageFormat;
use ray_tracing_in_one_weekend_compile_time::Color;

const IMAGE_WIDTH: usize = 256;
const IMAGE_HEIGHT: usize = 256;
const PIXEL_COUNT: usize = IMAGE_WIDTH * IMAGE_HEIGHT;

fn main() -> anyhow::Result<()> {
    const PIXEL_COLORS: [Color; PIXEL_COUNT] = ray_trace();
    const BUFFER: [u8; PIXEL_COUNT * 3] = pixels_to_buffer(PIXEL_COLORS);
    let img = image::RgbImage::from_raw(IMAGE_WIDTH as u32, IMAGE_HEIGHT as u32, BUFFER.to_vec())
        .unwrap();
    img.save_with_format("./result.png", ImageFormat::Png)?;

    Ok(())
}

const fn ray_trace() -> [Color; PIXEL_COUNT] {
    let mut pixel_colors = [Color::ZERO; PIXEL_COUNT];
    let mut i = 0;
    let mut j = (IMAGE_HEIGHT - 1) as i32;
    // forが使えないのでwhileで代用
    while 0 <= j {
        while i < IMAGE_WIDTH {
            pixel_colors[j as usize * IMAGE_WIDTH + i] = Color::new(
                (i as f64) / ((IMAGE_WIDTH - 1) as f64),
                j as f64 / (IMAGE_HEIGHT - 1) as f64,
                0.25,
            );
            i += 1;
        }
        i = 0;
        j -= 1;
    }
    pixel_colors
}

const fn pixels_to_buffer(pixels: [Color; PIXEL_COUNT]) -> [u8; PIXEL_COUNT * 3] {
    //let buf: Vec<_> = PIXEL_COLORS
    //    .iter()
    //    .map(|&(r, g, b)| (r * 255.999, g * 255.999, b * 255.999))
    //    .map(|&(r, g, b)| [r as u8, g as u8, b as u8])
    //    .flatten()
    //    .collect();
    //buf

    // 上と同じことをしている
    let mut buf = [0; PIXEL_COUNT * 3];
    let mut i = 0;
    while i < PIXEL_COUNT {
        buf[i * 3 + 0] = (pixels[i].x * 255.999) as u8;
        buf[i * 3 + 1] = (pixels[i].y * 255.999) as u8;
        buf[i * 3 + 2] = (pixels[i].z * 255.999) as u8;
        i += 1;
    }
    buf
}

まとめ

今回追加したコードもfaatureの使用は必要ですがほぼ通常通りの書き方で書けました。
abssqrtなんかの処理を自前で定義するのは面倒&不安なので標準のものをconst fnで使用できるようになって欲しいです。


次回はレイトレーシングっぽい処理が追加されるはずです。