1use futures::StreamExt;
2use gloo_net::websocket::Message;
3use gloo_net::websocket::futures::WebSocket;
4use leptos::html::Canvas;
5use leptos::logging::*;
6use leptos::prelude::*;
7use leptos::task::spawn_local;
8use leptos::wasm_bindgen::JsCast;
9use std::cell::RefCell;
10use std::rc::Rc;
11use web_sys::HtmlCanvasElement;
12use web_sys::WebGl2RenderingContext as GL;
13use web_sys::WebGlProgram;
14
15use crate::ArrayView;
16
17pub enum WaterfallMode {
18 Websocket(String),
19 Data(ReadSignal<Vec<u8>>),
20}
21
22impl Default for WaterfallMode {
23 fn default() -> Self {
24 Self::Websocket("ws://127.0.0.1:9001".to_string())
25 }
26}
27
28struct RenderState {
29 canvas: HtmlCanvasElement,
30 gl: GL,
31 shader: WebGlProgram,
32 texture_offset: i32,
33}
34
35const SHADER_HEIGHT: usize = 256;
36
37#[component]
38pub fn Waterfall(
40 #[prop(into)] min: Signal<f32>,
41 #[prop(into)] max: Signal<f32>,
42 #[prop(optional)] mode: WaterfallMode,
43) -> impl IntoView {
44 let data = match mode {
45 WaterfallMode::Data(d) => d,
46 WaterfallMode::Websocket(s) => {
47 let (data, set_data) = signal(vec![]);
48 spawn_local(async move {
49 let mut ws = WebSocket::open(&s).unwrap();
50 while let Some(msg) = ws.next().await {
51 match msg {
52 Ok(Message::Bytes(b)) => {
53 set_data.set(b);
54 }
55 _ => {
56 log!("TimeSink: WebSocket {:?}", msg);
57 }
58 }
59 }
60 log!("TimeSink: WebSocket Closed");
61 });
62 data
63 }
64 };
65
66 let canvas_ref = NodeRef::<Canvas>::new();
67 Effect::new(move || {
68 if let Some(canvas) = canvas_ref.get() {
69 let gl: GL = canvas
70 .get_context("webgl2")
71 .unwrap()
72 .unwrap()
73 .dyn_into()
74 .unwrap();
75
76 let vert_code = r"
77 attribute vec2 gTexCoord0;
78 varying vec2 coord;
79
80 void main()
81 {
82 gl_Position = vec4(gTexCoord0, 0, 1);
83 coord = gTexCoord0;
84 }
85 ";
86 let vert_shader = gl.create_shader(GL::VERTEX_SHADER).unwrap();
87 gl.shader_source(&vert_shader, vert_code);
88 gl.compile_shader(&vert_shader);
89
90 let frag_code = r"
91 precision mediump float;
92
93 varying vec2 coord;
94 uniform float u_min;
95 uniform float u_max;
96 uniform float yoffset;
97 uniform sampler2D frequency_data;
98
99 vec3 color_map(float t) {
100 const vec3 c0 = vec3(0.2777273272234177, 0.005407344544966578, 0.3340998053353061);
101 const vec3 c1 = vec3(0.1050930431085774, 1.404613529898575, 1.384590162594685);
102 const vec3 c2 = vec3(-0.3308618287255563, 0.214847559468213, 0.09509516302823659);
103 const vec3 c3 = vec3(-4.634230498983486, -5.799100973351585, -19.33244095627987);
104 const vec3 c4 = vec3(6.228269936347081, 14.17993336680509, 56.69055260068105);
105 const vec3 c5 = vec3(4.776384997670288, -13.74514537774601, -65.35303263337234);
106 const vec3 c6 = vec3(-5.435455855934631, 4.645852612178535, 26.3124352495832);
107
108 return c0+t*(c1+t*(c2+t*(c3+t*(c4+t*(c5+t*c6)))));
109 }
110
111 void main()
112 {
113 vec4 sample = texture2D(frequency_data, vec2(coord.x * 0.5 + 0.5, coord.y * 0.5 - 0.5 + yoffset));
114 float power = (10.0 * log(sample.r) / log(10.0) - u_min) / (u_max - u_min);
115 gl_FragColor = vec4(color_map(clamp(power, 0.0, 1.0)), 1.0);
116 }
117 ";
118
119 let frag_shader = gl.create_shader(GL::FRAGMENT_SHADER).unwrap();
120 gl.shader_source(&frag_shader, frag_code);
121 gl.compile_shader(&frag_shader);
122
123 let shader = gl.create_program().unwrap();
124 gl.attach_shader(&shader, &vert_shader);
125 gl.attach_shader(&shader, &frag_shader);
126 gl.link_program(&shader);
127 gl.use_program(Some(&shader));
128
129 let texture = gl.create_texture().unwrap();
130 gl.bind_texture(GL::TEXTURE_2D, Some(&texture));
131 gl.tex_parameteri(GL::TEXTURE_2D, GL::TEXTURE_WRAP_S, GL::REPEAT as i32);
132 gl.tex_parameteri(GL::TEXTURE_2D, GL::TEXTURE_WRAP_T, GL::REPEAT as i32);
133 gl.tex_parameteri(GL::TEXTURE_2D, GL::TEXTURE_MIN_FILTER, GL::NEAREST as i32);
134 gl.tex_parameteri(GL::TEXTURE_2D, GL::TEXTURE_MAG_FILTER, GL::NEAREST as i32);
135
136 let fft_size = 1024_usize;
137 initialize_texture(&gl, fft_size);
138
139 let vertexes = [-1.0, -1.0, -1.0, 1.0, 1.0, 1.0, 1.0, -1.0];
140 let vertex_buffer = gl.create_buffer().unwrap();
141 gl.bind_buffer(GL::ARRAY_BUFFER, Some(&vertex_buffer));
142 let view = unsafe { f32::view(&vertexes) };
143 gl.buffer_data_with_array_buffer_view(GL::ARRAY_BUFFER, &view, GL::STATIC_DRAW);
144
145 let indices = [0, 1, 2, 0, 2, 3];
146 let indices_buffer = gl.create_buffer().unwrap();
147 gl.bind_buffer(GL::ELEMENT_ARRAY_BUFFER, Some(&indices_buffer));
148 let view = unsafe { u16::view(&indices) };
149 gl.buffer_data_with_array_buffer_view(GL::ELEMENT_ARRAY_BUFFER, &view, GL::STATIC_DRAW);
150
151 let loc = gl.get_attrib_location(&shader, "gTexCoord0") as u32;
152 gl.enable_vertex_attrib_array(loc);
153 gl.vertex_attrib_pointer_with_i32(loc, 2, GL::FLOAT, false, 0, 0);
154
155 {
156 let gl = gl.clone();
157 let shader = shader.clone();
158 let _ = RenderEffect::new(move |_| {
159 let u_min = gl.get_uniform_location(&shader, "u_min");
160 gl.uniform1f(u_min.as_ref(), min.get());
161 let u_max = gl.get_uniform_location(&shader, "u_max");
162 gl.uniform1f(u_max.as_ref(), max.get());
163 });
164 }
165
166 let state = RenderState {
167 canvas,
168 gl,
169 shader,
170 texture_offset: 0,
171 };
172 request_animation_frame(render(Rc::new(RefCell::new(state)), data, fft_size))
173 }
174 });
175
176 view! { <canvas node_ref=canvas_ref style="width: 100%; height: 100%" /> }
177}
178
179fn render(
180 state: Rc<RefCell<RenderState>>,
181 data: ReadSignal<Vec<u8>>,
182 last_fft_size: usize,
183) -> impl FnOnce() + 'static {
184 move || {
185 let mut fft_size_val = last_fft_size;
186 {
187 let RenderState {
188 canvas,
189 gl,
190 shader,
191 texture_offset,
192 } = &mut (*state.borrow_mut());
193
194 let display_width = canvas.client_width() as u32;
195 let display_height = canvas.client_height() as u32;
196
197 let need_resize = canvas.width() != display_width || canvas.height() != display_height;
198
199 if need_resize {
200 canvas.set_width(display_width);
201 canvas.set_height(display_height);
202 gl.viewport(0, 0, display_width as i32, display_height as i32);
203 }
204
205 let bytes = &*data.read_untracked();
206 if !bytes.is_empty() {
207 fft_size_val = bytes.len() / 4;
208 if fft_size_val != last_fft_size {
209 initialize_texture(gl, fft_size_val);
210 }
211
212 let samples = unsafe {
213 let s = fft_size_val;
214 let p = bytes.as_ptr();
215 std::slice::from_raw_parts(p as *const f32, s)
216 };
217
218 let view = unsafe { f32::view(samples) };
219 gl.tex_sub_image_2d_with_i32_and_i32_and_u32_and_type_and_array_buffer_view_and_src_offset(
220 GL::TEXTURE_2D,
221 0,
222 0,
223 *texture_offset,
224 fft_size_val as i32,
225 1,
226 GL::RED,
227 GL::FLOAT,
228 &view,
229 0,
230 )
231 .unwrap();
232
233 let loc = gl.get_uniform_location(shader, "yoffset");
234 gl.uniform1f(loc.as_ref(), *texture_offset as f32 / SHADER_HEIGHT as f32);
235 *texture_offset = (*texture_offset + 1) % SHADER_HEIGHT as i32;
236 }
237
238 gl.draw_elements_with_i32(GL::TRIANGLES, 6, GL::UNSIGNED_SHORT, 0);
239 }
240 request_animation_frame(render(state, data, fft_size_val))
241 }
242}
243
244fn initialize_texture(gl: &GL, fft_size: usize) {
245 let texture = vec![0.0f32; fft_size * SHADER_HEIGHT];
246 let view = unsafe { f32::view(&texture) };
247 gl.tex_image_2d_with_i32_and_i32_and_i32_and_format_and_type_and_array_buffer_view_and_src_offset(
248 GL::TEXTURE_2D,
249 0,
250 GL::R32F as i32,
251 fft_size as i32,
252 SHADER_HEIGHT as i32,
253 0,
254 GL::RED,
255 GL::FLOAT,
256 &view,
257 0,
258 ).unwrap();
259}