dyna3/app-proto/src/display.rs

use core::array;
use nalgebra::{DMatrix, Rotation3, Vector3};
use sycamore::{prelude::*, motion::create_raf};
use web_sys::{
    console,
    window,
    KeyboardEvent,
    WebGl2RenderingContext,
    WebGlProgram,
    WebGlShader,
    WebGlUniformLocation,
    wasm_bindgen::{JsCast, JsValue}
};

use crate::AppState;

fn compile_shader(
    context: &WebGl2RenderingContext,
    shader_type: u32,
    source: &str,
) -> WebGlShader {
    let shader = context.create_shader(shader_type).unwrap();
    context.shader_source(&shader, source);
    context.compile_shader(&shader);
    shader
}

fn get_uniform_array_locations<const N: usize>(
    context: &WebGl2RenderingContext,
    program: &WebGlProgram,
    var_name: &str,
    member_name_opt: Option<&str>
) -> [Option<WebGlUniformLocation>; N] {
    array::from_fn(|n| {
        let name = match member_name_opt {
            Some(member_name) => format!("{var_name}[{n}].{member_name}"),
            None => format!("{var_name}[{n}]")
        };
        context.get_uniform_location(&program, name.as_str())
    })
}

// load the given data into the vertex input of the given name
fn bind_vertex_attrib(
    context: &WebGl2RenderingContext,
    index: u32,
    size: i32,
    data: &[f32]
) {
    // create a data buffer and bind it to ARRAY_BUFFER
    let buffer = context.create_buffer().unwrap();
    context.bind_buffer(WebGl2RenderingContext::ARRAY_BUFFER, Some(&buffer));
    
    // load the given data into the buffer. the function `Float32Array::view`
    // creates a raw view into our module's `WebAssembly.Memory` buffer.
    // allocating more memory will change the buffer, invalidating the view.
    // that means we have to make sure we don't allocate any memory until the
    // view is dropped
    unsafe {
        context.buffer_data_with_array_buffer_view(
            WebGl2RenderingContext::ARRAY_BUFFER,
            &js_sys::Float32Array::view(&data),
            WebGl2RenderingContext::STATIC_DRAW,
        );
    }
    
    // allow the target attribute to be used
    context.enable_vertex_attrib_array(index);
    
    // take whatever's bound to ARRAY_BUFFER---here, the data buffer created
    // above---and bind it to the target attribute
    //
    //   https://developer.mozilla.org/en-US/docs/Web/API/WebGLRenderingContext/vertexAttribPointer
    //
    context.vertex_attrib_pointer_with_i32(
        index,
        size,
        WebGl2RenderingContext::FLOAT,
        false, // don't normalize
        0, // zero stride
        0, // zero offset
    );
}

#[component]
pub fn Display() -> View {
    let state = use_context::<AppState>();
    
    // canvas
    let display = create_node_ref();
    
    // navigation
    let pitch_up = create_signal(0.0);
    let pitch_down = create_signal(0.0);
    let yaw_right = create_signal(0.0);
    let yaw_left = create_signal(0.0);
    let roll_ccw = create_signal(0.0);
    let roll_cw = create_signal(0.0);
    let zoom_in = create_signal(0.0);
    let zoom_out = create_signal(0.0);
    let turntable = create_signal(false); /* BENCHMARKING */
    
    // change listener
    let scene_changed = create_signal(true);
    create_effect(move || {
        state.assembly.elements.with(|elts| {
            for (_, elt) in elts {
                elt.representation.track();
            }
        });
        state.selection.track();
        scene_changed.set(true);
    });
    
    /* INSTRUMENTS */
    const SAMPLE_PERIOD: i32 = 60;
    let mut last_sample_time = 0.0;
    let mut frames_since_last_sample = 0;
    let mean_frame_interval = create_signal(0.0);
    
    on_mount(move || {
        // timing
        let mut last_time = 0.0;
        
        // viewpoint
        const ROT_SPEED: f64 = 0.4; // in radians per second
        const ZOOM_SPEED: f64 = 0.15; // multiplicative rate per second
        const TURNTABLE_SPEED: f64 = 0.1; /* BENCHMARKING */
        let mut orientation = DMatrix::<f64>::identity(5, 5);
        let mut rotation = DMatrix::<f64>::identity(5, 5);
        let mut location_z: f64 = 5.0;
        
        // display parameters
        const OPACITY: f32 = 0.5; /* SCAFFOLDING */
        const HIGHLIGHT: f32 = 0.2; /* SCAFFOLDING */
        const LAYER_THRESHOLD: i32 = 0; /* DEBUG */
        const DEBUG_MODE: i32 = 0; /* DEBUG */
        
        /* INSTRUMENTS */
        let performance = window().unwrap().performance().unwrap();
        
        // get the display canvas
        let canvas = display.get().unchecked_into::<web_sys::HtmlCanvasElement>();
        let ctx = canvas
            .get_context("webgl2")
            .unwrap()
            .unwrap()
            .dyn_into::<WebGl2RenderingContext>()
            .unwrap();
        
        // compile and attach the vertex and fragment shaders
        let vertex_shader = compile_shader(
            &ctx,
            WebGl2RenderingContext::VERTEX_SHADER,
            include_str!("identity.vert"),
        );
        let fragment_shader = compile_shader(
            &ctx,
            WebGl2RenderingContext::FRAGMENT_SHADER,
            include_str!("inversive.frag"),
        );
        let program = ctx.create_program().unwrap();
        ctx.attach_shader(&program, &vertex_shader);
        ctx.attach_shader(&program, &fragment_shader);
        ctx.link_program(&program);
        let link_status = ctx
            .get_program_parameter(&program, WebGl2RenderingContext::LINK_STATUS)
            .as_bool()
            .unwrap();
        let link_msg = if link_status {
            "Linked successfully"
        } else {
            "Linking failed"
        };
        console::log_1(&JsValue::from(link_msg));
        ctx.use_program(Some(&program));
        
        /* DEBUG */
        // print the maximum number of vectors that can be passed as
        // uniforms to a fragment shader. the OpenGL ES 3.0 standard
        // requires this maximum to be at least 224, as discussed in the
        // documentation of the GL_MAX_FRAGMENT_UNIFORM_VECTORS parameter
        // here:
        //
        //   https://registry.khronos.org/OpenGL-Refpages/es3.0/html/glGet.xhtml
        //
        // there are also other size limits. for example, on Aaron's
        // machine, the the length of a float or genType array seems to be
        // capped at 1024 elements
        console::log_2(
            &ctx.get_parameter(WebGl2RenderingContext::MAX_FRAGMENT_UNIFORM_VECTORS).unwrap(),
            &JsValue::from("uniform vectors available")
        );
        
        // find indices of vertex attributes and uniforms
        const SPHERE_MAX: usize = 200;
        let position_index = ctx.get_attrib_location(&program, "position") as u32;
        let sphere_cnt_loc = ctx.get_uniform_location(&program, "sphere_cnt");
        let sphere_sp_locs = get_uniform_array_locations::<SPHERE_MAX>(
            &ctx, &program, "sphere_list", Some("sp")
        );
        let sphere_lt_locs = get_uniform_array_locations::<SPHERE_MAX>(
            &ctx, &program, "sphere_list", Some("lt")
        );
        let color_locs = get_uniform_array_locations::<SPHERE_MAX>(
            &ctx, &program, "color_list", None
        );
        let highlight_locs = get_uniform_array_locations::<SPHERE_MAX>(
            &ctx, &program, "highlight_list", None
        );
        let resolution_loc = ctx.get_uniform_location(&program, "resolution");
        let shortdim_loc = ctx.get_uniform_location(&program, "shortdim");
        let opacity_loc = ctx.get_uniform_location(&program, "opacity");
        let layer_threshold_loc = ctx.get_uniform_location(&program, "layer_threshold");
        let debug_mode_loc = ctx.get_uniform_location(&program, "debug_mode");
        
        // create a vertex array and bind it to the graphics context
        let vertex_array = ctx.create_vertex_array().unwrap();
        ctx.bind_vertex_array(Some(&vertex_array));
        
        // set the vertex positions
        const VERTEX_CNT: usize = 6;
        let positions: [f32; 3*VERTEX_CNT] = [
            // northwest triangle
            -1.0, -1.0, 0.0,
            -1.0,  1.0, 0.0,
             1.0,  1.0, 0.0,
            // southeast triangle
            -1.0, -1.0, 0.0,
             1.0,  1.0, 0.0,
             1.0, -1.0, 0.0
        ];
        bind_vertex_attrib(&ctx, position_index, 3, &positions);
        
        // set up a repainting routine
        let (_, start_animation_loop, _) = create_raf(move || {
            // get the time step
            let time = performance.now();
            let time_step = 0.001*(time - last_time);
            last_time = time;
            
            // get the navigation state
            let pitch_up_val = pitch_up.get();
            let pitch_down_val = pitch_down.get();
            let yaw_right_val = yaw_right.get();
            let yaw_left_val = yaw_left.get();
            let roll_ccw_val = roll_ccw.get();
            let roll_cw_val = roll_cw.get();
            let zoom_in_val = zoom_in.get();
            let zoom_out_val = zoom_out.get();
            let turntable_val = turntable.get(); /* BENCHMARKING */
            
            // update the assembly's orientation
            let ang_vel = {
                let pitch = pitch_up_val - pitch_down_val;
                let yaw = yaw_right_val - yaw_left_val;
                let roll = roll_ccw_val - roll_cw_val;
                if pitch != 0.0 || yaw != 0.0 || roll != 0.0 {
                    ROT_SPEED * Vector3::new(-pitch, yaw, roll).normalize()
                } else {
                    Vector3::zeros()
                }
            } /* BENCHMARKING */ + if turntable_val {
                Vector3::new(0.0, TURNTABLE_SPEED, 0.0)
            } else {
                Vector3::zeros()
            };
            let mut rotation_sp = rotation.fixed_view_mut::<3, 3>(0, 0);
            rotation_sp.copy_from(
                Rotation3::from_scaled_axis(time_step * ang_vel).matrix()
            );
            orientation = &rotation * &orientation;
            
            // update the assembly's location
            let zoom = zoom_out_val - zoom_in_val;
            location_z *= (time_step * ZOOM_SPEED * zoom).exp();
            
            if scene_changed.get() {
                /* INSTRUMENTS */
                // measure mean frame interval
                frames_since_last_sample += 1;
                if frames_since_last_sample >= SAMPLE_PERIOD {
                    mean_frame_interval.set((time - last_sample_time) / (SAMPLE_PERIOD as f64));
                    last_sample_time = time;
                    frames_since_last_sample = 0;
                }
                
                // find the map from assembly space to world space
                let location = {
                    let u = -location_z;
                    DMatrix::from_column_slice(5, 5, &[
                        1.0, 0.0,   0.0, 0.0, 0.0,
                        0.0, 1.0,   0.0, 0.0, 0.0,
                        0.0, 0.0,   1.0, 0.0,   u,
                        0.0, 0.0, 2.0*u, 1.0, u*u,
                        0.0, 0.0,   0.0, 0.0, 1.0
                    ])
                };
                let assembly_to_world = &location * &orientation;
                
                // get the assembly
                let (
                    elt_cnt,
                    reps_world,
                    colors,
                    highlights
                ) = state.assembly.elements.with(|elts| {
                    (
                        // number of elements
                        elts.len() as i32,
                        
                        // representation vectors in world coordinates
                        elts.iter().map(
                            |(_, elt)| elt.representation.with(|rep| &assembly_to_world * rep)
                        ).collect::<Vec<_>>(),
                        
                        // colors
                        elts.iter().map(|(key, elt)| {
                            if state.selection.with(|sel| sel.contains(&key)) {
                                elt.color.map(|ch| 0.2 + 0.8*ch)
                            } else {
                                elt.color
                            }
                        }).collect::<Vec<_>>(),
                        
                        // highlight levels
                        elts.iter().map(|(key, _)| {
                            if state.selection.with(|sel| sel.contains(&key)) {
                                1.0_f32
                            } else {
                                HIGHLIGHT
                            }
                        }).collect::<Vec<_>>()
                    )
                });
                
                // set the resolution
                let width = canvas.width() as f32;
                let height = canvas.height() as f32;
                ctx.uniform2f(resolution_loc.as_ref(), width, height);
                ctx.uniform1f(shortdim_loc.as_ref(), width.min(height));
                
                // pass the assembly
                ctx.uniform1i(sphere_cnt_loc.as_ref(), elt_cnt);
                for n in 0..reps_world.len() {
                    let v = &reps_world[n];
                    ctx.uniform3f(
                        sphere_sp_locs[n].as_ref(),
                        v[0] as f32, v[1] as f32, v[2] as f32
                    );
                    ctx.uniform2f(
                        sphere_lt_locs[n].as_ref(),
                        v[3] as f32, v[4] as f32
                    );
                    ctx.uniform3fv_with_f32_array(
                        color_locs[n].as_ref(),
                        &colors[n]
                    );
                    ctx.uniform1f(
                        highlight_locs[n].as_ref(),
                        highlights[n]
                    );
                }
                
                // pass the display parameters
                ctx.uniform1f(opacity_loc.as_ref(), OPACITY);
                ctx.uniform1i(layer_threshold_loc.as_ref(), LAYER_THRESHOLD);
                ctx.uniform1i(debug_mode_loc.as_ref(), DEBUG_MODE);
                
                // draw the scene
                ctx.draw_arrays(WebGl2RenderingContext::TRIANGLES, 0, VERTEX_CNT as i32);
                
                // clear the scene change flag
                scene_changed.set(
                    pitch_up_val != 0.0
                    || pitch_down_val != 0.0
                    || yaw_left_val != 0.0
                    || yaw_right_val != 0.0
                    || roll_cw_val != 0.0
                    || roll_ccw_val != 0.0
                    || zoom_in_val != 0.0
                    || zoom_out_val != 0.0
                    || turntable_val /* BENCHMARKING */
                );
            } else {
                frames_since_last_sample = 0;
                mean_frame_interval.set(-1.0);
            }
        });
        start_animation_loop();
    });
    
    let set_nav_signal = move |event: KeyboardEvent, value: f64| {
        let mut navigating = true;
        let shift = event.shift_key();
        match event.key().as_str() {
            "ArrowUp" if shift => zoom_in.set(value),
            "ArrowDown" if shift => zoom_out.set(value),
            "ArrowUp" => pitch_up.set(value),
            "ArrowDown" => pitch_down.set(value),
            "ArrowRight" if shift => roll_cw.set(value),
            "ArrowLeft" if shift => roll_ccw.set(value),
            "ArrowRight" => yaw_right.set(value),
            "ArrowLeft" => yaw_left.set(value),
            _ => navigating = false
        };
        if navigating {
            scene_changed.set(true);
            event.prevent_default();
        }
    };
    
    view! {
        /* TO DO */
        // switch back to integer-valued parameters when that becomes possible
        // again
        canvas(
            ref=display,
            width="600",
            height="600",
            tabindex="0",
            on:keydown=move |event: KeyboardEvent| {
                if event.key() == "Shift" {
                    roll_cw.set(yaw_right.get());
                    roll_ccw.set(yaw_left.get());
                    zoom_in.set(pitch_up.get());
                    zoom_out.set(pitch_down.get());
                    yaw_right.set(0.0);
                    yaw_left.set(0.0);
                    pitch_up.set(0.0);
                    pitch_down.set(0.0);
                } else {
                    if event.key() == "Enter" { /* BENCHMARKING */
                        turntable.set_fn(|turn| !turn);
                        scene_changed.set(true);
                    }
                    set_nav_signal(event, 1.0);
                }
            },
            on:keyup=move |event: KeyboardEvent| {
                if event.key() == "Shift" {
                    yaw_right.set(roll_cw.get());
                    yaw_left.set(roll_ccw.get());
                    pitch_up.set(zoom_in.get());
                    pitch_down.set(zoom_out.get());
                    roll_cw.set(0.0);
                    roll_ccw.set(0.0);
                    zoom_in.set(0.0);
                    zoom_out.set(0.0);
                } else {
                    set_nav_signal(event, 0.0);
                }
            },
            on:blur=move |_| {
                pitch_up.set(0.0);
                pitch_down.set(0.0);
                yaw_right.set(0.0);
                yaw_left.set(0.0);
                roll_ccw.set(0.0);
                roll_cw.set(0.0);
            }
        )
    }
}