Skip to main content

mujoco_rs/
viewer.rs

1//! Module related to implementation of the [`MjViewer`]. For implementation of the C++ wrapper,
2//! see [`crate::cpp_viewer::MjViewerCpp`] (enabled by the `cpp-viewer` cargo feature).
3
4use std::sync::{Arc, Mutex, MutexGuard, PoisonError};
5use std::sync::atomic::{AtomicBool, Ordering};
6use std::time::{Duration, Instant};
7use std::ops::{Deref, DerefMut};
8use std::collections::BTreeSet;
9use std::num::NonZero;
10use std::error::Error;
11use std::fmt::Display;
12use std::borrow::Cow;
13
14use winit::event::{ElementState, KeyEvent, Modifiers, MouseButton, MouseScrollDelta, WindowEvent};
15#[cfg(feature = "viewer-ui")] use glutin::display::GetGlDisplay;
16use winit::platform::pump_events::EventLoopExtPumpEvents;
17use glutin::prelude::PossiblyCurrentGlContext;
18use winit::keyboard::{KeyCode, PhysicalKey};
19use winit::event_loop::EventLoop;
20use winit::dpi::PhysicalPosition;
21use glutin::surface::GlSurface;
22use winit::window::Fullscreen;
23use bitflags::bitflags;
24
25use crate::wrappers::mj_rendering::{MjrContext, MjrRectangle, MjtFont, MjtGridPos};
26use crate::util::{optional_sparse_addr_range, LockUnpoison, ThreeWayMerge};
27use crate::vis_common::{sync_geoms, flip_image_vertically, write_png};
28use crate::wrappers::mj_auxiliary::{MjVisual, MjStatistic};
29use crate::winit_gl_base::{RenderBaseGlState, RenderBase};
30use crate::wrappers::mj_primitive::{MjtNum, MjtSize};
31use crate::wrappers::mj_data::{MjData, MjtState};
32use crate::{builder_setters, mujoco_version};
33use crate::wrappers::mj_option::MjOption;
34use crate::wrappers::mj_visualization::*;
35use crate::wrappers::mj_editing::MjSpec;
36use crate::wrappers::mj_model::MjModel;
37
38
39#[cfg(feature = "viewer-ui")]
40mod ui;
41
42// Re-export egui for user convenience when using custom UI callbacks
43#[cfg(feature = "viewer-ui")]
44pub use egui;
45
46
47// Rust native viewer
48const MJ_VIEWER_DEFAULT_SIZE_PX: (u32, u32) = (1280, 720);
49const DOUBLE_CLICK_WINDOW_MS: u128 = 250;
50const TOUCH_BAR_ZOOM_FACTOR: f64 = 0.1;
51const FPS_SMOOTHING_FACTOR: f64 = 0.1;
52const REALTIME_FACTOR_SMOOTHING_FACTOR: f64 = 0.1;
53const REALTIME_FACTOR_DISPLAY_THRESHOLD: f64 = 0.02;
54
55/// How much extra room to create in the internal [`MjvScene`]. Useful for drawing labels, etc.
56pub(crate) const EXTRA_SCENE_GEOM_SPACE: usize = 2000;
57
58const HELP_MENU_TITLES: &str = concat!(
59    "Toggle help\n",
60    "Toggle info\n",
61    "Toggle v-sync\n",
62    "Toggle realtime check\n",
63    "Toggle full screen\n",
64    "Free camera\n",
65    "Track camera\n",
66    "Camera orbit\n",
67    "Camera pan\n",
68    "Camera look at\n",
69    "Zoom\n",
70    "Object select\n",
71    "Selection rotate\n",
72    "Selection translate\n",
73    "Exit\n",
74    "Reset simulation\n",
75    "Cycle cameras\n",
76    "Visualization toggles",
77);
78
79const HELP_MENU_VALUES: &str = concat!(
80    "F1\n",
81    "F2\n",
82    "F3\n",
83    "F4\n",
84    "F5\n",
85    "Escape\n",
86    "Control + Alt + double-left click\n",
87    "Left drag\n",
88    "Right [+Shift] drag\n",
89    "Alt + double-left click\n",
90    "Zoom, middle drag\n",
91    "Double-left click\n",
92    "Control + [Shift] + drag\n",
93    "Control + Alt + [Shift] + drag\n",
94    "Control + Q\n",
95    "Backspace\n",
96    "[ ]\n",
97    "See MjViewer docs"
98);
99
100// Compile-time check: HELP_MENU_TITLES and HELP_MENU_VALUES must have the same
101// number of newline-delimited entries, otherwise the help overlay misaligns.
102const fn count_byte(s: &[u8], target: u8) -> usize {
103    let mut count = 0;
104    let mut i = 0;
105    while i < s.len() {
106        if s[i] == target {
107            count += 1;
108        }
109        i += 1;
110    }
111    count
112}
113const _: () = assert!(
114    count_byte(HELP_MENU_TITLES.as_bytes(), b'\n')
115        == count_byte(HELP_MENU_VALUES.as_bytes(), b'\n'),
116    "HELP_MENU_TITLES and HELP_MENU_VALUES must have the same number of lines"
117);
118
119/// Errors that can occur when initializing or running the MuJoCo viewer.
120#[derive(Debug)]
121#[non_exhaustive]
122pub enum MjViewerError {
123    /// The event loop failed to initialize.
124    EventLoopError(winit::error::EventLoopError),
125    /// A glutin operation failed.
126    GlutinError(glutin::error::Error),
127    /// Returned when the egui painter (OpenGL UI renderer) fails to initialize.
128    PainterInitError(String),
129    /// OpenGL / window initialization failed.
130    GlInitFailed(crate::error::GlInitError),
131    /// A scene operation failed (e.g. user-scene sync overflowed the geom buffer).
132    SceneError(crate::error::MjSceneError),
133    /// A rendering-context operation failed (e.g. a pixel-read buffer is too small or the
134    /// viewport has invalid dimensions).
135    ContextError(crate::error::MjrContextError),
136    /// The model's structure signature does not match the viewer's passive model.
137    /// Call [`ViewerSharedState::sync_model`] or [`ViewerSharedState::sync_data`] first.
138    SignatureMismatch,
139    /// The asset ID is out of range.
140    IndexOutOfBounds {
141        /// The ID that was supplied.
142        id: usize,
143        /// The number of assets in the model (the valid range is `0..len`).
144        len: usize,
145    },
146}
147
148/// Formats a human-readable description of the viewer error.
149impl Display for MjViewerError {
150    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
151        match self {
152            Self::EventLoopError(e) => write!(f, "event loop failed to initialize: {e}"),
153            Self::GlutinError(e) => write!(f, "glutin error: {e}"),
154            Self::PainterInitError(e) => write!(f, "failed to initialize egui painter: {e}"),
155            Self::GlInitFailed(e) => write!(f, "GL initialization failed: {e}"),
156            Self::SceneError(e) => write!(f, "scene error: {e}"),
157            Self::ContextError(e) => write!(f, "rendering context error: {e}"),
158            Self::SignatureMismatch => write!(f, "model signature mismatch: call sync_model / sync_data first"),
159            Self::IndexOutOfBounds { id, len } => write!(f, "asset id {id} is out of range [0, {len})"),
160        }
161    }
162}
163
164/// Provides the underlying error source, if any.
165impl Error for MjViewerError {
166    fn source(&self) -> Option<&(dyn Error + 'static)> {
167        match self {
168            Self::EventLoopError(e) => Some(e),
169            Self::GlutinError(e) => Some(e),
170            Self::PainterInitError(_) => None,
171            Self::GlInitFailed(e) => Some(e),
172            Self::SceneError(e) => Some(e),
173            Self::ContextError(e) => Some(e),
174            Self::SignatureMismatch | Self::IndexOutOfBounds { .. } => None,
175        }
176    }
177}
178
179/// Converts an [`MjSceneError`](crate::error::MjSceneError) into [`MjViewerError::SceneError`].
180impl From<crate::error::MjSceneError> for MjViewerError {
181    fn from(e: crate::error::MjSceneError) -> Self {
182        Self::SceneError(e)
183    }
184}
185
186/// Converts a [`GlInitError`](crate::error::GlInitError) into [`MjViewerError::GlInitFailed`].
187impl From<crate::error::GlInitError> for MjViewerError {
188    fn from(e: crate::error::GlInitError) -> Self {
189        Self::GlInitFailed(e)
190    }
191}
192
193/// Converts an [`MjrContextError`](crate::error::MjrContextError) into [`MjViewerError::ContextError`].
194impl From<crate::error::MjrContextError> for MjViewerError {
195    fn from(e: crate::error::MjrContextError) -> Self {
196        Self::ContextError(e)
197    }
198}
199
200/// Converts a [`glutin::error::Error`] into [`MjViewerError::GlutinError`].
201impl From<glutin::error::Error> for MjViewerError {
202    fn from(e: glutin::error::Error) -> Self {
203        Self::GlutinError(e)
204    }
205}
206
207/// Internal state that is used by [`MjViewer`] to store
208/// [`MjData`]-related state. This is separate from [`MjViewer`]
209/// to allow use in multi-threaded programs, where the physics part
210/// runs in another thread and syncs the state with the viewer
211/// running in the main thread.
212/// 
213/// The state can be obtained through [`MjViewer::state`], which will return a reference to an `Arc<Mutex<ViewerSharedState>>`
214/// instance. For scoped access, you may also use [`MjViewer::with_state_lock`].
215#[derive(Debug)]
216pub struct ViewerSharedState {
217    /// This attribute, [`ViewerSharedState::data_passive`] and [`ViewerSharedState::data_passive_state_old`]
218    /// are used together to detect changes made to the state within the viewer.
219    /// This can happen due to changes made through the UI to joints, equalities, actuators, etc.
220    data_passive_state: Box<[MjtNum]>,
221    data_passive_state_old: Box<[MjtNum]>,
222    pub(crate) data_passive: MjData<Box<MjModel>>,
223    pert: MjvPerturb,
224    running: Arc<AtomicBool>,
225    user_scene: MjvScene,
226
227    /* Internals */
228    last_sync_time: Instant,
229    /// Time factor representing the ratio of viewer syncs with model's selected timestep.
230    realtime_factor_smooth: f64,
231    /// Preallocated buffer for storing the new [`MjData`] state.
232    data_state_buffer: Box<[MjtNum]>,
233    /// Timestamp when physics options (opt) were last synced
234    last_opt_sync_time: Instant,
235    /// Timestamp when visualization options (vis) were last synced
236    last_vis_sync_time: Instant,
237    /// Timestamp when statistics (stat) were last synced
238    last_stat_sync_time: Instant,
239
240    /* Model parameter change tracking */
241    prev_opt: MjOption,
242    prev_vis: MjVisual,
243    prev_stat: MjStatistic,
244
245    /* Pending GPU asset re-uploads: contains the IDs to upload on the next render call */
246    texture_reupload_pending: BTreeSet<usize>,
247    mesh_reupload_pending: BTreeSet<usize>,
248    hfield_reupload_pending: BTreeSet<usize>,
249}
250
251impl ViewerSharedState {
252    fn new<M: Deref<Target = MjModel>>(model: M, max_user_geom: usize) -> Self {
253        // Empty values to avoid unnecessary double creation
254        let empty: Box<[MjtNum]> = vec![].into_boxed_slice();
255        let empty_model = Box::new(MjSpec::new().compile().unwrap());
256        let empty_data = MjData::new(empty_model);
257        let empty_scene = MjvScene::new(empty_data.model(), 0);
258
259        let mut shared_state = Self {
260            data_passive: empty_data,
261            data_passive_state: empty.clone(),
262            data_passive_state_old: empty.clone(),
263            data_state_buffer: empty,
264            user_scene: empty_scene,
265            pert: MjvPerturb::default(),
266            running: Arc::new(AtomicBool::new(true)),
267            last_sync_time: Instant::now(),
268            realtime_factor_smooth: 1.0,
269            last_opt_sync_time: Instant::now(),
270            last_vis_sync_time: Instant::now(),
271            last_stat_sync_time: Instant::now(),
272            /* Model parameter change tracking */
273            prev_opt: MjOption::default(),
274            prev_vis: MjVisual::default(),
275            prev_stat: MjStatistic::default(),
276            /* Pending GPU asset re-uploads */
277            texture_reupload_pending: BTreeSet::new(),
278            mesh_reupload_pending: BTreeSet::new(),
279            hfield_reupload_pending: BTreeSet::new(),
280        };
281
282        shared_state.reload_model(&model, max_user_geom);
283        shared_state
284    }
285
286    /// Reinitializes all model-dependent internal state.
287    /// Called on construction and whenever [`_sync_data`](Self::_sync_data) or
288    /// [`sync_model`](Self::sync_model) detects a model change.
289    fn reload_model(&mut self, model: &MjModel, max_user_geom: usize) {
290        let model_passive = Box::new(model.clone());
291        self.data_passive = MjData::new(model_passive);
292        let model_passive = self.data_passive.model();
293
294        self.user_scene = MjvScene::new(model_passive, max_user_geom);
295        let state_size = model_passive.state_size(MjtState::mjSTATE_INTEGRATION as u32);
296        self.data_passive_state = vec![0.0; state_size].into_boxed_slice();
297        // Read the actual initial state (qpos0 may be non-zero) so that data_passive_state_old
298        // matches data_passive_state from the start, preventing a spurious write-back of the
299        // default pose to the incoming data on the first sync after a model change.
300        self.data_passive.read_state_into(
301            MjtState::mjSTATE_INTEGRATION as u32,
302            &mut self.data_passive_state,
303        );
304        self.data_passive_state_old = self.data_passive_state.clone();
305        self.data_state_buffer = self.data_passive_state.clone();
306        self.realtime_factor_smooth = 1.0;
307        self.pert = MjvPerturb::default();
308        // Clear pending re-uploads: the new context will already have the model's
309        // GPU resources loaded from scratch.
310        self.texture_reupload_pending.clear();
311        self.mesh_reupload_pending.clear();
312        self.hfield_reupload_pending.clear();
313    }
314
315    /// Checks whether the viewer is still running or is supposed to run.
316    pub fn running(&self) -> bool {
317        self.running.load(Ordering::Relaxed)
318    }
319
320    /// Returns an immutable reference to a user scene for drawing custom visual-only geoms.
321    /// Geoms in the user scene are preserved between calls to [`ViewerSharedState::sync_data`].
322    pub fn user_scene(&self) -> &MjvScene {
323        &self.user_scene
324    }
325
326    /// Returns a mutable reference to a user scene for drawing custom visual-only geoms.
327    /// Geoms in the user scene are preserved between calls to [`ViewerSharedState::sync_data`].
328    pub fn user_scene_mut(&mut self) -> &mut MjvScene {
329        &mut self.user_scene
330    }
331
332    /// Performs a bidirectional three-way merge of model's `opt`, `vis`, and `stat` parameters
333    /// between the viewer's passive model and the incoming model.
334    ///
335    /// Detects model changes via signature comparison and reloads internal state if needed.
336    pub fn sync_model(&mut self, model: &mut MjModel) {
337        // Check if model signature changed
338        if model.signature() != self.data_passive.model().signature() {
339            // Model changed: reload internal state, skip parameter sync
340            let max_user_geom = self.user_scene.maxgeom() as usize;
341            self.reload_model(model, max_user_geom);
342        }
343
344        self.sync_model_opt(model.opt_mut());
345        self.sync_model_vis(model.vis_mut());
346        self.sync_model_stat(model.stat_mut());
347    }
348
349    /// Performs a bidirectional three-way merge of [`MjModel::opt`] between the viewer's
350    /// passive model and the provided option struct.
351    ///
352    /// Changes made by the viewer UI are written to `opt`; changes made to `opt` externally
353    /// are written back to the viewer's passive model.
354    pub fn sync_model_opt(&mut self, opt: &mut MjOption) {
355        ThreeWayMerge::merge(opt, self.data_passive.model_opt_mut(), &mut self.prev_opt);
356        self.last_opt_sync_time = Instant::now();
357    }
358
359    /// Performs a bidirectional three-way merge of [`MjModel::vis`] between the viewer's
360    /// passive model and the provided visual struct.
361    ///
362    /// Changes made by the viewer UI are written to `vis`; changes made to `vis` externally
363    /// are written back to the viewer's passive model.
364    pub fn sync_model_vis(&mut self, vis: &mut MjVisual) {
365        ThreeWayMerge::merge(vis, self.data_passive.model_vis_mut(), &mut self.prev_vis);
366        self.last_vis_sync_time = Instant::now();
367    }
368
369    /// Performs a bidirectional three-way merge of [`MjModel::stat`] between the viewer's
370    /// passive model and the provided statistic struct.
371    ///
372    /// Changes made by the viewer UI are written to `stat`; changes made to `stat` externally
373    /// are written back to the viewer's passive model.
374    pub fn sync_model_stat(&mut self, stat: &mut MjStatistic) {
375        ThreeWayMerge::merge(stat, self.data_passive.model_stat_mut(), &mut self.prev_stat);
376        self.last_stat_sync_time = Instant::now();
377    }
378
379    /// Returns the last time physics options (opt) were synced.
380    pub fn last_opt_sync_time(&self) -> Instant {
381        self.last_opt_sync_time
382    }
383
384    /// Returns the last time visualization options (vis) were synced.
385    pub fn last_vis_sync_time(&self) -> Instant {
386        self.last_vis_sync_time
387    }
388
389    /// Returns the last time statistics (stat) were synced.
390    pub fn last_stat_sync_time(&self) -> Instant {
391        self.last_stat_sync_time
392    }
393
394    fn check_signature(&self, model: &MjModel) -> Result<(), MjViewerError> {
395        if model.signature() != self.data_passive.model().signature() {
396            return Err(MjViewerError::SignatureMismatch);
397        }
398        Ok(())
399    }
400
401    /// Copies the texture with `texture_id` from `model` into the viewer's internal passive model
402    /// and schedules a GPU re-upload for the next [`MjViewer::render`] call.
403    ///
404    /// The upload is processed on the next call to [`MjViewer::render`], at which point
405    /// the updated texture will be reflected in the scene.
406    ///
407    /// # Errors
408    /// - [`MjViewerError::SignatureMismatch`] if `model`'s signature does not match the viewer's passive model.
409    /// - [`MjViewerError::IndexOutOfBounds`] if `texture_id >= model.ntex()`.
410    pub fn update_texture_from(&mut self, model: &MjModel, texture_id: usize) -> Result<(), MjViewerError> {
411        self.check_signature(model)?;
412        let ntex = model.ntex() as usize;
413        if texture_id >= ntex {
414            return Err(MjViewerError::IndexOutOfBounds { id: texture_id, len: ntex });
415        }
416        let tex_adr = model.tex_adr()[texture_id] as usize;
417        let tex_width = model.tex_width()[texture_id] as usize;
418        let tex_height = model.tex_height()[texture_id] as usize;
419        let tex_nchannel = model.tex_nchannel()[texture_id] as usize;
420        let tex_len = tex_width * tex_height * tex_nchannel;
421        // SAFETY: Signature and bounds verified above; tex_len is derived from the texture's
422        // own metadata, so the slice cannot exceed the allocation.
423        unsafe { self.data_passive.model_mut() }
424            .tex_data_mut()[tex_adr..tex_adr + tex_len]
425            .copy_from_slice(&model.tex_data()[tex_adr..tex_adr + tex_len]);
426        self.texture_reupload_pending.insert(texture_id);
427        Ok(())
428    }
429
430    /// Copies all textures from `model` into the viewer's internal passive model
431    /// and schedules a GPU re-upload for the next [`MjViewer::render`] call.
432    ///
433    /// # Errors
434    /// - [`MjViewerError::SignatureMismatch`] if `model`'s signature does not match the viewer's passive model.
435    pub fn update_textures_from(&mut self, model: &MjModel) -> Result<(), MjViewerError> {
436        self.check_signature(model)?;
437        // SAFETY: Signature verified above, so tex_data layouts match exactly.
438        unsafe { self.data_passive.model_mut() }
439            .tex_data_mut()
440            .copy_from_slice(model.tex_data());
441        self.texture_reupload_pending.extend(0..model.ntex() as usize);
442        Ok(())
443    }
444
445    /// Copies the mesh with `mesh_id` from `model` into the viewer's internal passive model
446    /// and schedules a GPU re-upload for the next [`MjViewer::render`] call.
447    ///
448    /// All data arrays read by `mjr_uploadMesh` are copied: vertex positions
449    /// (`mesh_vert`), per-vertex normals (`mesh_normal`), UV texture coordinates
450    /// (`mesh_texcoord`), face--vertex indices (`mesh_face`), face--normal indices
451    /// (`mesh_facenormal`), face--texcoord indices (`mesh_facetexcoord`), and convex
452    /// hull graph data (`mesh_graph`). Layout fields (address and count arrays) are
453    /// not copied because they are fixed by the model signature.
454    ///
455    /// The upload is processed on the next call to [`MjViewer::render`], at which point
456    /// the updated mesh will be reflected in the scene.
457    ///
458    /// # Errors
459    /// - [`MjViewerError::SignatureMismatch`] if `model`'s signature does not match the viewer's passive model.
460    /// - [`MjViewerError::IndexOutOfBounds`] if `mesh_id >= model.nmesh()`.
461    pub fn update_mesh_from(&mut self, model: &MjModel, mesh_id: usize) -> Result<(), MjViewerError> {
462        self.check_signature(model)?;
463        let nmesh = model.nmesh() as usize;
464        if mesh_id >= nmesh {
465            return Err(MjViewerError::IndexOutOfBounds { id: mesh_id, len: nmesh });
466        }
467        let vert_adr = model.mesh_vertadr()[mesh_id] as usize;
468        let vert_num = model.mesh_vertnum()[mesh_id] as usize;
469        let normal_adr = model.mesh_normaladr()[mesh_id] as usize;
470        let normal_num = model.mesh_normalnum()[mesh_id] as usize;
471        let texcoord_adr = model.mesh_texcoordadr()[mesh_id];
472        let texcoord_num = model.mesh_texcoordnum()[mesh_id] as usize;
473        let face_adr = model.mesh_faceadr()[mesh_id] as usize;
474        let face_num = model.mesh_facenum()[mesh_id] as usize;
475        let graph_range = optional_sparse_addr_range(model.mesh_graphadr(), mesh_id, model.mesh_graph().len());
476        // SAFETY: Signature and bounds verified above; all offsets and counts are taken
477        // from the mesh's own metadata in a model with matching layout.
478        let passive = unsafe { self.data_passive.model_mut() };
479        passive.mesh_vert_mut()[vert_adr..vert_adr + vert_num]
480            .copy_from_slice(&model.mesh_vert()[vert_adr..vert_adr + vert_num]);
481        passive.mesh_normal_mut()[normal_adr..normal_adr + normal_num]
482            .copy_from_slice(&model.mesh_normal()[normal_adr..normal_adr + normal_num]);
483        if texcoord_adr >= 0 {
484            let texcoord_adr = texcoord_adr as usize;
485            passive.mesh_texcoord_mut()[texcoord_adr..texcoord_adr + texcoord_num]
486                .copy_from_slice(&model.mesh_texcoord()[texcoord_adr..texcoord_adr + texcoord_num]);
487        }
488        unsafe {
489            passive.mesh_face_mut()[face_adr..face_adr + face_num]
490                .copy_from_slice(&model.mesh_face()[face_adr..face_adr + face_num]);
491            passive.mesh_facenormal_mut()[face_adr..face_adr + face_num]
492                .copy_from_slice(&model.mesh_facenormal()[face_adr..face_adr + face_num]);
493            passive.mesh_facetexcoord_mut()[face_adr..face_adr + face_num]
494                .copy_from_slice(&model.mesh_facetexcoord()[face_adr..face_adr + face_num]);
495            if let Some((graph_adr, mesh_graph_len)) = graph_range {
496                passive.mesh_graph_mut()[graph_adr..graph_adr + mesh_graph_len]
497                    .copy_from_slice(&model.mesh_graph()[graph_adr..graph_adr + mesh_graph_len]);
498            }
499        }
500        self.mesh_reupload_pending.insert(mesh_id);
501        Ok(())
502    }
503
504    /// Copies all meshes from `model` into the viewer's internal passive model
505    /// and schedules a GPU re-upload for the next [`MjViewer::render`] call.
506    ///
507    /// All data arrays read by `mjr_uploadMesh` are bulk-copied: vertex positions
508    /// (`mesh_vert`), per-vertex normals (`mesh_normal`), UV texture coordinates
509    /// (`mesh_texcoord`), face--vertex indices (`mesh_face`), face--normal indices
510    /// (`mesh_facenormal`), face--texcoord indices (`mesh_facetexcoord`), and convex
511    /// hull graph data (`mesh_graph`). Layout fields (address and count arrays) are
512    /// not copied because they are fixed by the model signature.
513    ///
514    /// # Errors
515    /// - [`MjViewerError::SignatureMismatch`] if `model`'s signature does not match the viewer's passive model.
516    pub fn update_meshes_from(&mut self, model: &MjModel) -> Result<(), MjViewerError> {
517        self.check_signature(model)?;
518        // SAFETY: Signature verified above, ensuring all mesh array layouts match exactly.
519        let passive = unsafe { self.data_passive.model_mut() };
520        passive.mesh_vert_mut().copy_from_slice(model.mesh_vert());
521        passive.mesh_normal_mut().copy_from_slice(model.mesh_normal());
522        passive.mesh_texcoord_mut().copy_from_slice(model.mesh_texcoord());
523        unsafe {
524            passive.mesh_face_mut().copy_from_slice(model.mesh_face());
525            passive.mesh_facenormal_mut().copy_from_slice(model.mesh_facenormal());
526            passive.mesh_facetexcoord_mut().copy_from_slice(model.mesh_facetexcoord());
527            passive.mesh_graph_mut().copy_from_slice(model.mesh_graph());
528        }
529        self.mesh_reupload_pending.extend(0..model.nmesh() as usize);
530        Ok(())
531    }
532
533    /// Copies the heightfield with `hfield_id` from `model` into the viewer's internal passive model
534    /// and schedules a GPU re-upload for the next [`MjViewer::render`] call.
535    ///
536    /// The upload is processed on the next call to [`MjViewer::render`], at which point
537    /// the updated heightfield will be reflected in the scene.
538    ///
539    /// # Errors
540    /// - [`MjViewerError::SignatureMismatch`] if `model`'s signature does not match the viewer's passive model.
541    /// - [`MjViewerError::IndexOutOfBounds`] if `hfield_id >= model.nhfield()`.
542    pub fn update_hfield_from(&mut self, model: &MjModel, hfield_id: usize) -> Result<(), MjViewerError> {
543        self.check_signature(model)?;
544        let nhfield = model.nhfield() as usize;
545        if hfield_id >= nhfield {
546            return Err(MjViewerError::IndexOutOfBounds { id: hfield_id, len: nhfield });
547        }
548        let hfield_adr = model.hfield_adr()[hfield_id] as usize;
549        let hfield_nrow = model.hfield_nrow()[hfield_id] as usize;
550        let hfield_ncol = model.hfield_ncol()[hfield_id] as usize;
551        let hfield_len = hfield_nrow * hfield_ncol;
552        // SAFETY: Signature and bounds verified above; hfield_len is derived from the
553        // heightfield's own row/column metadata.
554        unsafe { self.data_passive.model_mut() }
555            .hfield_data_mut()[hfield_adr..hfield_adr + hfield_len]
556            .copy_from_slice(&model.hfield_data()[hfield_adr..hfield_adr + hfield_len]);
557        self.hfield_reupload_pending.insert(hfield_id);
558        Ok(())
559    }
560
561    /// Copies all heightfields from `model` into the viewer's internal passive model
562    /// and schedules a GPU re-upload for the next [`MjViewer::render`] call.
563    ///
564    /// # Errors
565    /// - [`MjViewerError::SignatureMismatch`] if `model`'s signature does not match the viewer's passive model.
566    pub fn update_hfields_from(&mut self, model: &MjModel) -> Result<(), MjViewerError> {
567        self.check_signature(model)?;
568        // SAFETY: Signature verified above, so hfield_data layouts match exactly.
569        unsafe { self.data_passive.model_mut() }
570            .hfield_data_mut()
571            .copy_from_slice(model.hfield_data());
572        self.hfield_reupload_pending.extend(0..model.nhfield() as usize);
573        Ok(())
574    }
575
576
577    /// Same as [`ViewerSharedState::sync_data`], except it copies the entire [`MjData`]
578    /// struct (including large Jacobian and other arrays), not just the state needed for visualization.
579    ///
580    /// # Panics
581    /// Panics if the internal data copy fails due to an inconsistent model state (indicates a bug).
582    pub fn sync_data_full<M: Deref<Target = MjModel>>(&mut self, data: &mut MjData<M>) {
583        self._sync_data(data, true);
584    }
585
586    /// Syncs the viewer's internal passive [`MjData`] with `data`.
587    /// Synchronization happens in three steps:
588    ///
589    /// 1. The viewer checks if any changes have been made to the internal [`MjData`]
590    ///    since the last call to this method (since the last sync). Changed elements
591    ///    are selectively merged into `data` (elements the viewer did not touch are
592    ///    preserved).
593    /// 2. `data` is copied into the viewer's internal passive copy
594    ///    (visualization fields only; see warning below).
595    /// 3. Perturbations are applied to `data` via [`MjvPerturb::apply`], which
596    ///    **unconditionally zeroes `xfrc_applied`** before writing any active
597    ///    perturbation forces. Any external forces previously set on `data` will be
598    ///    cleared.
599    ///
600    /// Note that users must afterward call [`MjViewer::render`] for the scene
601    /// to be rendered and the UI to be processed.
602    ///
603    /// <div class="warning">
604    /// The user's data is copied into the viewer's internal passive copy via ``mjv_copyData``,
605    /// which skips large computed arrays not required for visualization.
606    /// The viewer's passive copy will therefore **not** contain:
607    ///
608    /// - mass matrices (``qM``, ``qLD``, ``qLDiagInv``, ``qLU``);
609    /// - constraint arrays (``efc_*``, ``iefc_*``, including constraint Jacobians).
610    ///
611    /// In UI callbacks these fields will be absent unless
612    /// [`ViewerSharedState::sync_data_full`] is used or they are recomputed explicitly
613    /// (e.g. via `data.forward()`).
614    ///
615    /// Additionally, because the viewer may write integration state (e.g. ``ctrl``) back
616    /// to the user's `data`, any Jacobians or other derived quantities in `data` may be
617    /// stale after this call and should be recomputed if needed.
618    /// </div>
619    ///
620    /// # Panics
621    /// Panics if the internal data copy or state merge fails due to an inconsistent model
622    /// state (indicates a bug).
623    pub fn sync_data<M: Deref<Target = MjModel>>(&mut self, data: &mut MjData<M>) {
624        self._sync_data(data, false);
625    }
626
627    /// Data sync implementation.
628    fn _sync_data<M: Deref<Target = MjModel>>(&mut self, data: &mut MjData<M>, full_sync: bool) {
629        // Recreate internal data and user scene when the model changes
630        if data.model().signature() != self.data_passive.model().signature() {
631            let max_user_geom = self.user_scene.maxgeom() as usize;
632            self.reload_model(data.model(), max_user_geom);
633        }
634
635        // Update statistics
636        let passive_time = self.data_passive.time();
637        let active_time = data.time();
638        if passive_time > 0.0 && active_time > passive_time {  // time = 0 means data was reset
639            let time_elapsed_sim = active_time - passive_time;
640            let elapsed_sync = self.last_sync_time.elapsed();
641            if !elapsed_sync.is_zero() {
642                self.realtime_factor_smooth += REALTIME_FACTOR_SMOOTHING_FACTOR * (
643                    time_elapsed_sim / elapsed_sync.as_secs_f64()
644                    - self.realtime_factor_smooth
645                );
646            }
647        } else {  // simulation was reset
648            self.realtime_factor_smooth = 1.0;
649        }
650
651        self.last_sync_time = Instant::now();
652
653        // Sync
654        self.data_passive.read_state_into(
655            MjtState::mjSTATE_INTEGRATION as u32,
656            &mut self.data_passive_state
657        );
658        if self.data_passive_state != self.data_passive_state_old {
659            data.read_state_into(MjtState::mjSTATE_INTEGRATION as u32, &mut self.data_state_buffer);
660            for ((new, passive), passive_old) in self.data_state_buffer.iter_mut()
661                .zip(&mut self.data_passive_state)
662                .zip(&mut self.data_passive_state_old)
663            {
664                if *passive_old != *passive {
665                    *new = *passive;
666                }
667            }
668
669            data.set_state(&self.data_state_buffer, MjtState::mjSTATE_INTEGRATION as u32)
670                .unwrap();
671        }
672
673        if full_sync {
674            // Copy everything.
675            data.copy_to(&mut self.data_passive)
676                .unwrap();
677        } else {
678            // Copy only visually-required information to the internal passive data.
679            data.copy_visual_to(&mut self.data_passive)
680                .unwrap();
681        }
682
683        // Make both saved states the same.
684        // If any modification is made through the viewer
685        // between syncs, then the above if block will trigger a transfer.
686        self.data_passive.read_state_into(  // read to match the synced passive MjData
687            MjtState::mjSTATE_INTEGRATION as u32,
688            &mut self.data_passive_state
689        );
690        self.data_passive_state_old.copy_from_slice(&self.data_passive_state);
691
692        // Apply perturbations
693        self.pert.apply(data);
694    }
695}
696
697
698/// A Rust-native implementation of the MuJoCo viewer. To conform to Rust's safety rules,
699/// the viewer doesn't store a mutable reference to the [`MjData`] struct, but it instead
700/// accepts it as a parameter in its methods.
701/// 
702/// The [`MjViewer::sync_data`] method must be called to sync the state of [`MjViewer`] and [`MjData`].
703/// 
704/// # Shortcuts
705/// Main keyboard and mouse shortcuts can be viewed by pressing `F1`.
706/// Additionally, some visualization toggles are included, but not displayed
707/// in the `F1` help menu:
708/// - C: camera,
709/// - U: actuator,
710/// - J: joint,
711/// - M: center of mass,
712/// - H: convex hull,
713/// - Z: light,
714/// - T: transparent,
715/// - I: inertia,
716/// - E: constraint.
717/// 
718/// # Safety
719/// Due to the nature of OpenGL, this should only be run in the **main thread**.
720#[derive(Debug)]
721pub struct MjViewer {
722    /* MuJoCo rendering */
723    scene: MjvScene,
724    context: MjrContext,
725    camera: MjvCamera,
726
727    /* Other MuJoCo related */
728    opt: MjvOption,
729
730    /* Internal state */
731    last_x: f64,
732    last_y: f64,
733    last_bnt_press_time: Instant,
734    rect_view: MjrRectangle,
735    rect_full: MjrRectangle,
736    fps_timer: Instant,
737    fps_smooth: f64,
738
739    /* OpenGL */
740    adapter: RenderBase,
741    event_loop: EventLoop<()>,
742    modifiers: Modifiers,
743    buttons_pressed: ButtonsPressed,
744    raw_cursor_position: (f64, f64),
745
746    /* External interaction */
747    shared_state: Arc<Mutex<ViewerSharedState>>,
748    /// Shared atomic flag cloned from [`ViewerSharedState::running`].
749    /// Allows [`MjViewer::running`] to read the flag without locking the mutex on every call
750    /// (the user's sim loop polls this on every iteration).
751    running_flag: Arc<AtomicBool>,
752
753    /* User interface */
754    #[cfg(feature = "viewer-ui")]
755    ui: ui::ViewerUI,
756
757    status: ViewerStatusBit,
758
759    /// Cached number of MJCF-defined cameras (`model.ncam`).
760    /// Updated in [`update_scene`](Self::update_scene) whenever the model changes.
761    /// Avoids locking `shared_state` just to read this integer in [`cycle_camera`](Self::cycle_camera).
762    ncam: MjtSize,
763
764    /// Pending screenshot request. [`Some`] with `(viewport_only, depth)` flags when a
765    /// screenshot is queued; [`None`] otherwise.
766    screenshot_pending: Option<(bool, bool)>
767}
768
769impl MjViewer {
770    /// Launches the MuJoCo viewer. A [`Result`] struct is returned that either contains
771    /// [`MjViewer`] or a [`MjViewerError`]. The `max_user_geom` parameter
772    /// defines how much space will be allocated for additional, user-defined visual-only geoms.
773    /// It can thus be set to 0 if no additional geoms will be drawn by the user.
774    ///
775    /// Note that the use of [`MjViewerBuilder`] is preferred, because it is more flexible.
776    /// Call [`MjViewer::builder`] to create a [`MjViewerBuilder`] instance.
777    /// # Returns
778    /// On success, returns [`Ok`] variant containing the [`MjViewer`].
779    /// # Errors
780    /// - [`MjViewerError::EventLoopError`] if the event loop fails to initialize.
781    /// - [`MjViewerError::GlInitFailed`] if OpenGL / window initialization fails.
782    /// - [`MjViewerError::GlutinError`] if a glutin operation fails.
783    /// - [`MjViewerError::PainterInitError`] if the UI painter fails to initialize
784    ///   (feature `viewer-ui`).
785    pub fn launch_passive<M: Deref<Target = MjModel>>(model: M, max_user_geom: usize) -> Result<Self, MjViewerError> {
786        MjViewerBuilder::new()
787            .max_user_geoms(max_user_geom)
788            .build_passive(model)
789    }
790
791    /// A shortcut for creating an instance of [`MjViewerBuilder`].
792    /// The builder can be used to build the viewer after configuring it.
793    /// It allows better configuration than [`MjViewer::launch_passive`], which
794    /// is fixed to achieve backward compatibility.
795    pub fn builder() -> MjViewerBuilder {
796        MjViewerBuilder::new()
797    }
798
799    /// Checks whether the viewer is still running.
800    pub fn running(&self) -> bool {
801        self.running_flag.load(Ordering::Relaxed)
802    }
803
804    /// Returns a reference to the shared state [`ViewerSharedState`].
805    /// This struct can be used to sync the state of the viewer with
806    /// the simulation, possibly running in another thread.
807    pub fn state(&self) -> &Arc<Mutex<ViewerSharedState>> {
808        &self.shared_state
809    }
810
811    /// Acquires a Mutex lock on the [`MjViewer`]'s shared state ([`MjViewer::state`]).
812    /// The acquired lock is passed to the function/closure `fun`.
813    ///
814    /// # Errors
815    /// Returns [`PoisonError`] if the mutex holding the shared state has panicked, thus poisoning
816    /// the lock.
817    ///
818    /// # Example
819    /// ```no_run
820    /// # use mujoco_rs::viewer::MjViewer;
821    /// # use mujoco_rs::prelude::*;
822    /// # let model = MjModel::from_xml_string("<mujoco/>").unwrap();
823    /// let mut viewer = MjViewer::builder().max_user_geoms(1)
824    ///     .build_passive(&model).unwrap();
825    /// viewer.with_state_lock(|mut lock| {
826    ///     let scene = lock.user_scene_mut();
827    ///     scene.create_geom(MjtGeom::mjGEOM_BOX, Some([1.0, 1.0, 1.0]), Some([0.0, 0.0, 0.0]), None, None);
828    /// }).unwrap();
829    /// ```
830    pub fn with_state_lock<F, R>(&self, fun: F) -> Result<R, PoisonError<MutexGuard<'_, ViewerSharedState>>>
831        where F: FnOnce(MutexGuard<ViewerSharedState>) -> R
832    {
833        Ok(fun(self.shared_state.lock()?))
834    }
835
836    /// Adds a user-defined UI callback for custom widgets in the viewer's UI.
837    /// The callback receives an [`egui::Context`] reference and can be used to create
838    /// custom windows, panels, or other UI elements.
839    /// It also receives a mutable reference to [`MjData`], which can be used to read
840    /// and modify simulation state. Note that the model can be accessed through [`MjData::model`].
841    ///
842    /// # Note
843    /// The viewer's internal shared-state [`Mutex`] is **held for the entire
844    /// duration of the callback** (because `data` is a live borrow of the guarded
845    /// `data_passive` field). Do **not** attempt to lock the shared state again from
846    /// within the callback as that will deadlock the viewer thread.
847    ///
848    /// # Example
849    /// ```no_run
850    /// # use mujoco_rs::prelude::*;
851    /// # use mujoco_rs::viewer::MjViewer;
852    /// # let model = MjModel::from_xml_string("<mujoco/>").unwrap();
853    /// # let mut viewer = MjViewer::launch_passive(&model, 0).unwrap();
854    /// viewer.add_ui_callback(|ctx, data| {
855    ///     use mujoco_rs::viewer::egui;
856    ///     egui::Window::new("Custom controls")
857    ///         .scroll(true)
858    ///         .show(ctx, |ui| {
859    ///             ui.label("Custom UI element");
860    ///         });
861    /// });
862    /// ```
863    #[cfg(feature = "viewer-ui")]
864    pub fn add_ui_callback<F>(&mut self, callback: F)
865    where
866        F: FnMut(&egui::Context, &mut MjData<Box<MjModel>>) + 'static
867    {
868        self.ui.add_ui_callback(callback);
869    }
870
871    /// Same as [`MjViewer::add_ui_callback`], except the `callback` does
872    /// not receive the passive [`MjData`] instance of the viewer.
873    /// Consequently, the mutex of the viewer's shared state doesn't need to
874    /// be locked, yielding better performance.
875    #[cfg(feature = "viewer-ui")]
876    pub fn add_ui_callback_detached<F>(&mut self, callback: F)
877    where
878        F: FnMut(&egui::Context) + 'static
879    {
880        self.ui.add_ui_callback_detached(callback);
881    }
882
883    /// Same as [`MjViewer::sync_data`], except it copies the entire [`MjData`]
884    /// struct (including large Jacobian and other arrays), not just the state needed for visualization.
885    /// This is a proxy to [`ViewerSharedState::sync_data_full`].
886    ///
887    /// # Panics
888    /// Panics if the internal data copy fails due to an inconsistent model state (indicates a bug).
889    pub fn sync_data_full<M: Deref<Target = MjModel>>(&mut self, data: &mut MjData<M>) {
890        self.shared_state.lock_unpoison().sync_data_full(data);
891    }
892
893    /// Syncs the state of viewer's internal [`MjData`] with `data`.
894    /// This is a proxy to [`ViewerSharedState::sync_data`].
895    /// 
896    /// Any changes made to the internal [`MjData`] in between syncs
897    /// get selectively merged back into `data` before the copy.
898    /// Perturbations are applied to `data` **after** the sync, which
899    /// **unconditionally zeroes `xfrc_applied`** (see
900    /// [`ViewerSharedState::sync_data`] for details).
901    /// 
902    /// Note that users must afterward call [`MjViewer::render`] for the scene
903    /// to be rendered and the UI to be processed.
904    /// 
905    /// <div class="warning">
906    /// The user's data is copied into the viewer's internal passive copy via ``mjv_copyData``,
907    /// which skips large computed arrays not required for visualization.
908    /// The viewer's passive copy will therefore **not** contain:
909    ///
910    /// - mass matrices (``qM``, ``qLD``, ``qLDiagInv``, ``qLU``);
911    /// - constraint arrays (``efc_*``, ``iefc_*``, including constraint Jacobians).
912    ///
913    /// In UI callbacks these fields will be absent unless
914    /// [`MjViewer::sync_data_full`] is used or they are recomputed explicitly
915    /// (e.g. via `data.forward()`).
916    ///
917    /// Additionally, because the viewer may write integration state (e.g. ``ctrl``) back
918    /// to the user's `data`, any Jacobians or other derived quantities in `data` may be
919    /// stale after this call and should be recomputed if needed.
920    /// </div>
921    ///
922    /// # Panics
923    /// Panics if the internal data copy or state merge fails due to an inconsistent model
924    /// state (indicates a bug).
925    ///
926    /// # Example
927    /// ```no_run
928    /// # use mujoco_rs::prelude::*;
929    /// # use mujoco_rs::viewer::MjViewer;
930    /// # let model = MjModel::from_xml("/path/scene.xml").unwrap();
931    /// # let mut viewer = MjViewer::builder().build_passive(&model).unwrap();
932    /// # let mut data = MjData::new(&model);
933    /// viewer.sync_data(&mut data);  // sync the data
934    /// viewer.render().unwrap();  // render the scene and process the user interface
935    /// ```
936    pub fn sync_data<M: Deref<Target = MjModel>>(&mut self, data: &mut MjData<M>) {
937        self.shared_state.lock_unpoison().sync_data(data);
938    }
939
940    /// Performs a bidirectional three-way merge of model parameters between the
941    /// viewer's passive model and the incoming model.
942    /// This is a proxy to [`ViewerSharedState::sync_model`].
943    ///
944    /// Detects model changes via signature comparison and reloads internal state if needed.
945    /// After a reload, the passive model mirrors the incoming model's defaults, so the
946    /// first merge is effectively a no-op. On subsequent calls, viewer UI changes and
947    /// simulation-side changes are merged bidirectionally.
948    pub fn sync_model(&mut self, model: &mut MjModel) {
949        self.shared_state.lock_unpoison().sync_model(model);
950    }
951
952    /// Performs a bidirectional three-way merge of [`MjModel::opt`] between the viewer's
953    /// passive model and the provided option struct.
954    /// This is a proxy to [`ViewerSharedState::sync_model_opt`].
955    ///
956    /// This allows updating physics options without requiring `unsafe` access via
957    /// [`MjData::model_mut`] or manipulating the viewer's shared state directly.
958    pub fn sync_model_opt(&mut self, opt: &mut MjOption) {
959        self.shared_state.lock_unpoison().sync_model_opt(opt);
960    }
961
962    /// Performs a bidirectional three-way merge of [`MjModel::vis`] between the viewer's
963    /// passive model and the provided visual struct.
964    /// This is a proxy to [`ViewerSharedState::sync_model_vis`].
965    ///
966    /// This allows updating visualization options without requiring `unsafe` access via
967    /// [`MjData::model_mut`] or manipulating the viewer's shared state directly.
968    pub fn sync_model_vis(&mut self, vis: &mut MjVisual) {
969        self.shared_state.lock_unpoison().sync_model_vis(vis);
970    }
971
972    /// Performs a bidirectional three-way merge of [`MjModel::stat`] between the viewer's
973    /// passive model and the provided statistic struct.
974    /// This is a proxy to [`ViewerSharedState::sync_model_stat`].
975    ///
976    /// This allows updating model statistics without requiring `unsafe` access via
977    /// [`MjData::model_mut`] or manipulating the viewer's shared state directly.
978    pub fn sync_model_stat(&mut self, stat: &mut MjStatistic) {
979        self.shared_state.lock_unpoison().sync_model_stat(stat);
980    }
981
982    /// Copies the texture with `texture_id` from `model` into the viewer's internal passive model
983    /// and schedules a GPU re-upload for the next [`render`](Self::render) call.
984    ///
985    /// This is a proxy to [`ViewerSharedState::update_texture_from`].
986    ///
987    /// # Errors
988    /// - [`MjViewerError::SignatureMismatch`] if `model`'s signature does not match the viewer's passive model.
989    /// - [`MjViewerError::IndexOutOfBounds`] if `texture_id >= model.ntex()`.
990    pub fn update_texture_from(&mut self, model: &MjModel, texture_id: usize) -> Result<(), MjViewerError> {
991        self.shared_state.lock_unpoison().update_texture_from(model, texture_id)
992    }
993
994    /// Copies all textures from `model` into the viewer's internal passive model
995    /// and schedules a GPU re-upload for the next [`render`](Self::render) call.
996    ///
997    /// This is a proxy to [`ViewerSharedState::update_textures_from`].
998    ///
999    /// # Errors
1000    /// - [`MjViewerError::SignatureMismatch`] if `model`'s signature does not match the viewer's passive model.
1001    pub fn update_textures_from(&mut self, model: &MjModel) -> Result<(), MjViewerError> {
1002        self.shared_state.lock_unpoison().update_textures_from(model)
1003    }
1004
1005    /// Copies the mesh with `mesh_id` from `model` into the viewer's internal passive model
1006    /// and schedules a GPU re-upload for the next [`render`](Self::render) call.
1007    ///
1008    /// This is a proxy to [`ViewerSharedState::update_mesh_from`].
1009    ///
1010    /// # Errors
1011    /// - [`MjViewerError::SignatureMismatch`] if `model`'s signature does not match the viewer's passive model.
1012    /// - [`MjViewerError::IndexOutOfBounds`] if `mesh_id >= model.nmesh()`.
1013    pub fn update_mesh_from(&mut self, model: &MjModel, mesh_id: usize) -> Result<(), MjViewerError> {
1014        self.shared_state.lock_unpoison().update_mesh_from(model, mesh_id)
1015    }
1016
1017    /// Copies all meshes from `model` into the viewer's internal passive model
1018    /// and schedules a GPU re-upload for the next [`render`](Self::render) call.
1019    ///
1020    /// This is a proxy to [`ViewerSharedState::update_meshes_from`].
1021    ///
1022    /// # Errors
1023    /// - [`MjViewerError::SignatureMismatch`] if `model`'s signature does not match the viewer's passive model.
1024    pub fn update_meshes_from(&mut self, model: &MjModel) -> Result<(), MjViewerError> {
1025        self.shared_state.lock_unpoison().update_meshes_from(model)
1026    }
1027
1028    /// Copies the heightfield with `hfield_id` from `model` into the viewer's internal passive model
1029    /// and schedules a GPU re-upload for the next [`render`](Self::render) call.
1030    ///
1031    /// This is a proxy to [`ViewerSharedState::update_hfield_from`].
1032    ///
1033    /// # Errors
1034    /// - [`MjViewerError::SignatureMismatch`] if `model`'s signature does not match the viewer's passive model.
1035    /// - [`MjViewerError::IndexOutOfBounds`] if `hfield_id >= model.nhfield()`.
1036    pub fn update_hfield_from(&mut self, model: &MjModel, hfield_id: usize) -> Result<(), MjViewerError> {
1037        self.shared_state.lock_unpoison().update_hfield_from(model, hfield_id)
1038    }
1039
1040    /// Copies all heightfields from `model` into the viewer's internal passive model
1041    /// and schedules a GPU re-upload for the next [`render`](Self::render) call.
1042    ///
1043    /// This is a proxy to [`ViewerSharedState::update_hfields_from`].
1044    ///
1045    /// # Errors
1046    /// - [`MjViewerError::SignatureMismatch`] if `model`'s signature does not match the viewer's passive model.
1047    pub fn update_hfields_from(&mut self, model: &MjModel) -> Result<(), MjViewerError> {
1048        self.shared_state.lock_unpoison().update_hfields_from(model)
1049    }
1050
1051
1052    /// Processes the UI (when enabled), processes events, draws the scene
1053    /// and swaps buffers in OpenGL.
1054    ///
1055    /// # Errors
1056    /// - [`MjViewerError::GlutinError`] if the OpenGL context cannot be made current or the buffer swap fails.
1057    /// - [`MjViewerError::SceneError`] if synchronizing user scene geoms fails (e.g. the scene is
1058    ///   full).
1059    /// - [`MjViewerError::ContextError`] if reading pixels for a pending screenshot fails.
1060    pub fn render(&mut self) -> Result<(), MjViewerError> {
1061        let RenderBaseGlState {
1062            gl_context,
1063            gl_surface,
1064            ..
1065        } = self.adapter.state.as_ref().unwrap();
1066
1067        // Make sure everything is done on the viewer's window
1068        gl_context.make_current(gl_surface)?;
1069
1070        // Read the screen size
1071        self.update_rectangles(self.adapter.state.as_ref().unwrap().window.inner_size().into());
1072
1073        // Process mouse and keyboard events
1074        self.process_events();
1075
1076        // Update the scene from data and render
1077        self.update_scene()?;
1078
1079        // Viewport-only screenshot: capture the centered scene before the UI
1080        // panel is drawn on top (the scene is rendered to rect_full, so it
1081        // fills the entire window).
1082        if let Some((true, _)) = self.screenshot_pending {
1083            let (_, depth) = self.screenshot_pending.take().unwrap();
1084            self.capture_screenshot(depth)?;
1085        }
1086
1087        // Draw the user menu on top
1088        #[cfg(feature = "viewer-ui")]
1089        self.process_user_ui();
1090
1091        // Update the user menu state and overlays
1092        self.update_menus();
1093
1094        // Full-window screenshot: capture after all rendering (UI + overlays).
1095        if let Some((false, _)) = self.screenshot_pending {
1096            let (_, depth) = self.screenshot_pending.take().unwrap();
1097            self.capture_screenshot(depth)?;
1098        }
1099
1100        // Flush to the GPU
1101        self.swap_buffers()
1102    }
1103
1104    /// Perform OpenGL buffer swap.
1105    fn swap_buffers(&self) -> Result<(), MjViewerError> {
1106        let RenderBaseGlState {
1107            gl_context,
1108            gl_surface,
1109            ..
1110        } = self.adapter.state.as_ref().unwrap();
1111
1112        // Swap OpenGL buffers (render to screen)
1113        gl_surface.swap_buffers(gl_context).map_err(MjViewerError::GlutinError)
1114    }
1115
1116    /// Captures the current framebuffer contents as a PNG screenshot.
1117    /// Always reads from [`rect_full`](Self::rect_full) (the entire window).
1118    /// The caller controls what is visible by choosing *when* to call this
1119    /// method in the render pipeline. When `depth` is `true`, a 16-bit
1120    /// grayscale depth image is saved instead of an RGB image.
1121    ///
1122    /// # Errors
1123    /// Returns [`MjViewerError::ContextError`] if reading pixels from the framebuffer fails.
1124    fn capture_screenshot(&self, depth: bool) -> Result<(), MjViewerError> {
1125        let rect = &self.rect_full;
1126
1127        let w = rect.width as usize;
1128        let h = rect.height as usize;
1129        if w == 0 || h == 0 {
1130            return Ok(());
1131        }
1132
1133        // Generate a timestamped filename.
1134        let timestamp = std::time::SystemTime::now()
1135            .duration_since(std::time::UNIX_EPOCH)
1136            .unwrap_or_default()
1137            .as_secs();
1138
1139        if depth {
1140            let mut depth_buf = vec![0.0f32; w * h];
1141            self.context.read_pixels(None, Some(&mut depth_buf), rect)
1142                .map_err(MjViewerError::ContextError)?;
1143
1144            // OpenGL reads bottom-up; flip for top-down PNG row order.
1145            flip_image_vertically(&mut depth_buf, h, w);
1146
1147            // Linearize raw OpenGL depth into metric distance.
1148            let (map, stat) = {
1149                let lock = self.shared_state.lock_unpoison();
1150                let model = lock.data_passive.model();
1151                (model.vis().map.clone(), model.stat().clone())
1152            };
1153            let extent = stat.extent as f32;
1154            let near = map.znear * extent;
1155            let far = map.zfar * extent;
1156            for value in &mut depth_buf {
1157                *value = near / (1.0 - *value * (1.0 - near / far));
1158            }
1159
1160            // Normalize to 16-bit range.
1161            let max = depth_buf.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
1162            let min = depth_buf.iter().cloned().fold(f32::INFINITY, f32::min);
1163            let scale = u16::MAX as f32;
1164            let range = max - min;
1165            let encoded: Box<[u8]> = depth_buf
1166                .iter()
1167                .flat_map(|&x| {
1168                    let norm = if range > 0.0 { (x - min) / range } else { 0.0 };
1169                    ((norm * scale).clamp(0.0, scale) as u16).to_be_bytes()
1170                })
1171                .collect();
1172
1173            let path = format!("screenshot_{timestamp}_depth.png");
1174            if let Err(e) = write_png(
1175                &path, &encoded, w as u32, h as u32,
1176                png::ColorType::Grayscale, png::BitDepth::Sixteen, png::Compression::High
1177            ) {
1178                eprintln!("depth screenshot failed: {e}");
1179            } else {
1180                eprintln!("depth screenshot saved to {path} (min={min:.4}, max={max:.4})");
1181            }
1182        } else {
1183            let mut rgb = vec![0u8; w * h * 3];
1184            self.context.read_pixels(Some(&mut rgb), None, rect)
1185                .map_err(MjViewerError::ContextError)?;
1186
1187            // OpenGL reads bottom-up; flip for top-down PNG row order.
1188            flip_image_vertically(&mut rgb, h, w * 3);
1189
1190            let path = format!("screenshot_{timestamp}.png");
1191            if let Err(e) = write_png(
1192                &path, &rgb, w as u32, h as u32,
1193                png::ColorType::Rgb, png::BitDepth::Eight, png::Compression::High
1194            ) {
1195                eprintln!("screenshot failed: {e}");
1196            } else {
1197                eprintln!("screenshot saved to {path}");
1198            }
1199        }
1200        Ok(())
1201    }
1202
1203    fn update_smooth_fps(&mut self) {
1204        let elapsed = self.fps_timer.elapsed();
1205
1206        let fps = if elapsed.is_zero() {
1207            self.fps_smooth
1208        } else {
1209            1.0 / elapsed.as_secs_f64()
1210        };
1211
1212        self.fps_timer = Instant::now();
1213        self.fps_smooth += FPS_SMOOTHING_FACTOR * (fps - self.fps_smooth);
1214    }
1215
1216    /// Updates the scene and draws it to the display.
1217    fn update_scene(&mut self) -> Result<(), MjViewerError> {
1218        {
1219            let mut lock = self.shared_state.lock_unpoison();
1220            let ViewerSharedState {
1221                data_passive, pert, user_scene,
1222                texture_reupload_pending, mesh_reupload_pending, hfield_reupload_pending,
1223                ..
1224            } = lock.deref_mut();
1225
1226            // Recreate scene when the model changes
1227            if data_passive.model().signature() != self.scene.signature() {
1228                let new_model = data_passive.model();
1229                let ngeom = new_model.ffi().ngeom as usize;
1230                let max_user_geom = user_scene.maxgeom() as usize;
1231                self.scene = MjvScene::new(
1232                    new_model,
1233                    ngeom + max_user_geom + EXTRA_SCENE_GEOM_SPACE
1234                );
1235
1236                // Reset to a free camera: a tracking or fixed camera may reference a body
1237                // or camera ID that does not exist in the new model.
1238                self.camera = MjvCamera::new_free(new_model);
1239                // Recreate the rendering context so that GPU resources (textures,
1240                // meshes, heightfields, skins) match the new model.
1241                // SAFETY: the GL context was made current in render() before this call.
1242                self.context = unsafe { MjrContext::new(new_model) };
1243                self.ncam = new_model.ffi().ncam;
1244                #[cfg(feature = "viewer-ui")]
1245                self.ui.update_names(new_model);
1246                // reload_model already cleared the pending flags and notified waiters.
1247            }
1248
1249            // Process any pending GPU asset re-uploads. The GL context is current here
1250            // (ensured by render()). Each set is taken (leaving an empty set) and iterated.
1251            for id in std::mem::take(texture_reupload_pending) {
1252                self.context.upload_texture(data_passive.model(), id)?;
1253            }
1254            for id in std::mem::take(mesh_reupload_pending) {
1255                self.context.upload_mesh(data_passive.model(), id)?;
1256            }
1257            for id in std::mem::take(hfield_reupload_pending) {
1258                self.context.upload_hfield(data_passive.model(), id)?;
1259            }
1260
1261            // Update and render the scene from the MjData state
1262            self.scene.update(data_passive, &self.opt, pert, &mut self.camera);
1263
1264            // Draw geoms drawn through the user scene.
1265            sync_geoms(user_scene, &mut self.scene)?;
1266        }
1267        self.scene.render(&self.rect_full, &self.context);
1268        Ok(())
1269    }
1270
1271    /// Draws the user menu
1272    fn update_menus(&mut self) {
1273        let rectangle_from_ui = self.rect_view;
1274        let rectangle_full = self.rect_full;
1275
1276        // Overlay section
1277        if self.status.contains(ViewerStatusBit::HELP) {  // Help
1278            self.context.overlay(
1279                MjtFont::mjFONT_NORMAL, MjtGridPos::mjGRID_TOPLEFT,
1280                rectangle_from_ui,
1281                HELP_MENU_TITLES,
1282                Some(HELP_MENU_VALUES)
1283            );
1284        }
1285
1286        // Read later-required information and then drop the mutex lock
1287        let (
1288            time,
1289            memory_pct,
1290            mut total_memory,
1291            realtime_factor
1292        ) = {
1293            let state_lock = self.shared_state.lock_unpoison();
1294            let data_lock = &state_lock.data_passive;
1295            let memory_total = data_lock.narena().max(1) as f64;
1296            (
1297                data_lock.time(),
1298                100.0 * data_lock.maxuse_arena() as f64 / memory_total, memory_total,
1299                state_lock.realtime_factor_smooth
1300            )
1301        };
1302
1303        if self.status.contains(ViewerStatusBit::INFO) {  // Info
1304            self.update_smooth_fps();
1305
1306            // Overlay headers
1307            let headers = concat!(
1308                "FPS\n",
1309                "Time\n",
1310                "Memory\n",
1311                "Realtime factor"
1312            );
1313
1314            // Calculate the amount of memory used and represent with SI units.
1315            let mut memory_unit = ' ';
1316            if total_memory > 1e6 {
1317                total_memory /= 1e6;
1318                memory_unit = 'M';
1319            } else if total_memory > 1e3 {
1320                total_memory /= 1e3;
1321                memory_unit = 'k';
1322            }
1323
1324            // Format values of the overlay.
1325            let values = format!(
1326                concat!(
1327                    "{:.1}\n",
1328                    "{:.1}\n",
1329                    "{:.1} % out of {:.1} {}\n",
1330                    "{:.1} %"
1331                ),
1332                self.fps_smooth,
1333                time,
1334                memory_pct, total_memory, memory_unit,
1335                realtime_factor * 100.0
1336            );
1337
1338            self.context.overlay(
1339                MjtFont::mjFONT_NORMAL,
1340                MjtGridPos::mjGRID_BOTTOMLEFT,
1341                rectangle_from_ui,
1342                headers,
1343                Some(&values)
1344            );
1345        }
1346
1347        // Check for slowdowns
1348        if self.status.contains(ViewerStatusBit::WARN_REALTIME)
1349            && (realtime_factor - 1.0).abs() > REALTIME_FACTOR_DISPLAY_THRESHOLD
1350        {
1351            self.context.overlay(
1352                MjtFont::mjFONT_BIG,
1353                MjtGridPos::mjGRID_BOTTOMRIGHT,
1354                rectangle_full,
1355                &format!("Realtime factor: {:.1} %", realtime_factor * 100.0),
1356                None
1357            );
1358        }
1359    }
1360
1361    /// Gains scoped access to [`egui::Context`], which is part of the UI,
1362    /// for dealing with custom initialization (e.g., loading in images).
1363    #[cfg(feature = "viewer-ui")]
1364    pub fn with_ui_egui_ctx<F>(&mut self, once_fn: F)
1365        where F: FnOnce(&egui::Context)
1366    {
1367        self.ui.with_egui_ctx(once_fn);
1368    }
1369
1370    /// Draws the user UI
1371    #[cfg(feature = "viewer-ui")]
1372    fn process_user_ui(&mut self) {
1373        // Draw the user interface
1374
1375        use crate::viewer::ui::UiEvent;
1376        let RenderBaseGlState {window, ..} = &self.adapter.state.as_ref().unwrap();
1377
1378        let inner_size = window.inner_size();
1379        self.ui.init_2d();
1380        let left = self.ui.process(
1381            window, &mut self.status,
1382            &mut self.scene, &mut self.opt,
1383            &mut self.camera, &self.shared_state,
1384        );
1385
1386        // Adjust the viewport so MuJoCo doesn't draw over the UI
1387        self.rect_view.left = left as i32;
1388        self.rect_view.width = inner_size.width as i32;
1389
1390        // Reset some OpenGL settings so that MuJoCo can still draw
1391        self.ui.reset();
1392
1393        // Process events made in the user UI
1394        while let Some(event) = self.ui.drain_events() {
1395            use UiEvent::*;
1396            match event {
1397                Close => self.running_flag.store(false, Ordering::Relaxed),
1398                Fullscreen => self.toggle_full_screen(),
1399                ResetSimulation => {
1400                    let mut lock = self.shared_state.lock_unpoison();
1401                    lock.data_passive.reset();
1402                    lock.data_passive.forward();
1403                },
1404                AlignCamera => {
1405                    self.camera = MjvCamera::new_free(self.shared_state.lock_unpoison().data_passive.model());
1406                },
1407                VSyncToggle => {
1408                    self.update_vsync();
1409                },
1410                Screenshot { viewport_only, depth } => {
1411                    self.screenshot_pending = Some((viewport_only, depth));
1412                }
1413            }
1414        }
1415    }
1416
1417    /// Reads the state of requested vsync setting and makes appropriate calls to [`glutin`].
1418    fn update_vsync(&mut self) {
1419        let RenderBaseGlState {
1420            gl_surface, gl_context, ..
1421        } = &self.adapter.state.as_ref().unwrap();
1422
1423        if self.status.contains(ViewerStatusBit::VSYNC) {
1424            if let Err(e) = gl_surface.set_swap_interval(
1425                gl_context, glutin::surface::SwapInterval::Wait(NonZero::<u32>::MIN)
1426            ) {
1427                eprintln!("failed to enable vsync: {e}");
1428                self.status.set(ViewerStatusBit::VSYNC, false);
1429            }
1430        } else {
1431            if let Err(e) = gl_surface.set_swap_interval(
1432                gl_context, glutin::surface::SwapInterval::DontWait
1433            ) {
1434                eprintln!("failed to disable vsync: {e}");
1435                self.status.set(ViewerStatusBit::VSYNC, true);
1436            }
1437        }
1438    }
1439
1440    /// Updates the dimensions of the rectangles defining the dimensions of
1441    /// the user interface, as well as the actual scene viewer.
1442    fn update_rectangles(&mut self, viewport_size: (i32, i32)) {
1443        // The scene (middle) rectangle
1444        self.rect_view.width = viewport_size.0;
1445        self.rect_view.height = viewport_size.1;
1446
1447        self.rect_full.width = viewport_size.0;
1448        self.rect_full.height = viewport_size.1;
1449    }
1450
1451    /// Processes user input events.
1452    fn process_events(&mut self) {
1453        self.event_loop.pump_app_events(Some(Duration::ZERO), &mut self.adapter);
1454        while let Some(window_event) = self.adapter.queue.pop_front() {
1455            #[cfg(feature = "viewer-ui")]
1456            {
1457                let window: &winit::window::Window = &self.adapter.state.as_ref().unwrap().window;
1458                self.ui.handle_events(window, &window_event);
1459
1460                // if the UI has an active input focus, ignore all keyboard events
1461                if let WindowEvent::KeyboardInput { .. } = &window_event
1462                    && self.ui.focused()
1463                {
1464                    continue;
1465                }
1466            }
1467
1468            match window_event {
1469                WindowEvent::ModifiersChanged(modifiers) => self.modifiers = modifiers,
1470                WindowEvent::MouseInput {state, button, .. } => {
1471                    let is_pressed = state == ElementState::Pressed;
1472                    
1473                    #[cfg(feature = "viewer-ui")]
1474                    if self.ui.covered() && is_pressed {
1475                        continue;
1476                    }
1477
1478                    let index = match button {
1479                        MouseButton::Left => {
1480                            self.process_left_click(state);
1481                            ButtonsPressed::LEFT
1482                        },
1483                        MouseButton::Middle => ButtonsPressed::MIDDLE,
1484                        MouseButton::Right => ButtonsPressed::RIGHT,
1485                        _ => continue
1486                    };
1487
1488                    self.buttons_pressed.set(index, is_pressed);
1489                }
1490
1491                WindowEvent::CursorMoved { position, .. } => {
1492                    let PhysicalPosition { x, y } = position;
1493
1494                    // The UI might not be detected as covered as dragging can happen slightly outside
1495                    // of a (popup) window. This might seem like an ad-hoc solution, but is at the time the
1496                    // shortest and most efficient one.
1497                    #[cfg(feature = "viewer-ui")]
1498                    if self.ui.dragged() {
1499                        continue;
1500                    }
1501
1502                    self.process_cursor_pos(x, y);
1503                }
1504
1505                // Set the viewer's state to pending exit.
1506                WindowEvent::KeyboardInput {
1507                    event: KeyEvent {
1508                        physical_key: PhysicalKey::Code(KeyCode::KeyQ),
1509                        state: ElementState::Pressed, ..
1510                    }, ..
1511                } if self.modifiers.state().control_key()  => {
1512                    self.running_flag.store(false, Ordering::Relaxed);
1513                }
1514
1515                // Also set the viewer's state to pending exit if the window no longer exists.
1516                WindowEvent::CloseRequested => { self.running_flag.store(false, Ordering::Relaxed) }
1517
1518                // Free the camera from tracking.
1519                WindowEvent::KeyboardInput {
1520                    event: KeyEvent {
1521                        physical_key: PhysicalKey::Code(KeyCode::Escape),
1522                        state: ElementState::Pressed, ..
1523                    }, ..
1524                } => {
1525                    self.camera.free();
1526                }
1527
1528                // Toggle help menu
1529                WindowEvent::KeyboardInput {
1530                    event: KeyEvent {
1531                        physical_key: PhysicalKey::Code(KeyCode::F1),
1532                        state: ElementState::Pressed, ..
1533                    }, ..
1534                } => {
1535                    self.status.toggle(ViewerStatusBit::HELP);
1536                }
1537
1538                 // Toggle info menu
1539                WindowEvent::KeyboardInput {
1540                    event: KeyEvent {
1541                        physical_key: PhysicalKey::Code(KeyCode::F2),
1542                        state: ElementState::Pressed, ..
1543                    }, ..
1544                } => {
1545                    self.status.toggle(ViewerStatusBit::INFO);
1546                }
1547
1548                // Toggle VSync
1549                WindowEvent::KeyboardInput {
1550                    event: KeyEvent {physical_key: PhysicalKey::Code(KeyCode::F3), state: ElementState::Pressed, ..},
1551                    ..
1552                } => {
1553                    self.status.toggle(ViewerStatusBit::VSYNC);
1554                    self.update_vsync();
1555                }
1556
1557                // Non-realtime warnings
1558                WindowEvent::KeyboardInput {
1559                    event: KeyEvent {physical_key: PhysicalKey::Code(KeyCode::F4), state: ElementState::Pressed, ..},
1560                    ..
1561                } => {
1562                    self.status.toggle(ViewerStatusBit::WARN_REALTIME);
1563                }
1564
1565                // Full screen
1566                WindowEvent::KeyboardInput {
1567                    event: KeyEvent {
1568                        physical_key: PhysicalKey::Code(KeyCode::F5),
1569                        state: ElementState::Pressed, ..
1570                    }, ..
1571                } => {
1572                    self.toggle_full_screen();
1573                }
1574
1575                // Reset the simulation (the data).
1576                WindowEvent::KeyboardInput {
1577                    event: KeyEvent {
1578                        physical_key: PhysicalKey::Code(KeyCode::Backspace),
1579                        state: ElementState::Pressed, ..
1580                    }, ..
1581                } => {
1582                    let mut lock = self.shared_state.lock_unpoison();
1583                    lock.data_passive.reset();
1584                    lock.data_passive.forward();
1585                }
1586
1587                // Cycle to the next camera
1588                WindowEvent::KeyboardInput {
1589                    event: KeyEvent {
1590                        physical_key: PhysicalKey::Code(KeyCode::BracketRight),
1591                        state: ElementState::Pressed, ..
1592                    }, ..
1593                } => {
1594                    self.cycle_camera(1);
1595                }
1596
1597                // Cycle to the previous camera
1598                WindowEvent::KeyboardInput {
1599                    event: KeyEvent {
1600                        physical_key: PhysicalKey::Code(KeyCode::BracketLeft),
1601                        state: ElementState::Pressed, ..
1602                    }, ..
1603                } => {
1604                    self.cycle_camera(-1);
1605                }
1606
1607                // Toggles camera visualization.
1608                WindowEvent::KeyboardInput {
1609                    event: KeyEvent {physical_key: PhysicalKey::Code(KeyCode::KeyC), state: ElementState::Pressed, ..},
1610                    ..
1611                } => {
1612                    if self.modifiers.state().control_key() {  // Control + C is reserved for copy
1613                        continue;
1614                    }
1615
1616                    self.toggle_opt_flag(MjtVisFlag::mjVIS_CAMERA);
1617                }
1618
1619                WindowEvent::KeyboardInput {
1620                    event: KeyEvent {physical_key: PhysicalKey::Code(KeyCode::KeyU), state: ElementState::Pressed, ..},
1621                    ..
1622                } => self.toggle_opt_flag(MjtVisFlag::mjVIS_ACTUATOR),
1623
1624                WindowEvent::KeyboardInput {
1625                    event: KeyEvent {physical_key: PhysicalKey::Code(KeyCode::KeyJ), state: ElementState::Pressed, ..},
1626                    ..
1627                } => self.toggle_opt_flag(MjtVisFlag::mjVIS_JOINT),
1628
1629                WindowEvent::KeyboardInput {
1630                    event: KeyEvent {physical_key: PhysicalKey::Code(KeyCode::KeyM), state: ElementState::Pressed, ..},
1631                    ..
1632                } => self.toggle_opt_flag(MjtVisFlag::mjVIS_COM),
1633
1634                WindowEvent::KeyboardInput {
1635                    event: KeyEvent {physical_key: PhysicalKey::Code(KeyCode::KeyH), state: ElementState::Pressed, ..},
1636                    ..
1637                } => self.toggle_opt_flag(MjtVisFlag::mjVIS_CONVEXHULL),
1638
1639                WindowEvent::KeyboardInput {
1640                    event: KeyEvent {physical_key: PhysicalKey::Code(KeyCode::KeyZ), state: ElementState::Pressed, ..},
1641                    ..
1642                } => self.toggle_opt_flag(MjtVisFlag::mjVIS_LIGHT),
1643
1644                WindowEvent::KeyboardInput {
1645                    event: KeyEvent {physical_key: PhysicalKey::Code(KeyCode::KeyT), state: ElementState::Pressed, ..},
1646                    ..
1647                } => self.toggle_opt_flag(MjtVisFlag::mjVIS_TRANSPARENT),
1648
1649                WindowEvent::KeyboardInput {
1650                    event: KeyEvent {physical_key: PhysicalKey::Code(KeyCode::KeyI), state: ElementState::Pressed, ..},
1651                    ..
1652                } => self.toggle_opt_flag(MjtVisFlag::mjVIS_INERTIA),
1653
1654                WindowEvent::KeyboardInput {
1655                    event: KeyEvent {physical_key: PhysicalKey::Code(KeyCode::KeyE), state: ElementState::Pressed, ..},
1656                    ..
1657                } => self.toggle_opt_flag(MjtVisFlag::mjVIS_CONSTRAINT),
1658
1659                #[cfg(feature = "viewer-ui")]
1660                WindowEvent::KeyboardInput {
1661                    event: KeyEvent {physical_key: PhysicalKey::Code(KeyCode::KeyX), state: ElementState::Pressed, ..},
1662                    ..
1663                } => self.status.toggle(ViewerStatusBit::UI),
1664
1665                // Zoom in/out
1666                WindowEvent::MouseWheel {delta, ..} => {
1667                    #[cfg(feature = "viewer-ui")]
1668                    if self.ui.covered() {
1669                        continue;
1670                    }
1671
1672                    let value = match delta {
1673                        MouseScrollDelta::LineDelta(_, down) => down as f64,
1674                        MouseScrollDelta::PixelDelta(PhysicalPosition {y, ..}) => y * TOUCH_BAR_ZOOM_FACTOR
1675                    };
1676                    self.process_scroll(value);
1677                }
1678
1679                _ => {}  // ignore other events
1680            }
1681        }
1682    }
1683
1684    /// Toggles visualization options.
1685    fn toggle_opt_flag(&mut self, flag: MjtVisFlag) {
1686        let index = flag as usize;
1687        self.opt.flags[index] = 1 - self.opt.flags[index];
1688    }
1689
1690    /// Cycle MJCF defined cameras.
1691    fn cycle_camera(&mut self, direction: i32) {
1692        let ncam = self.ncam;
1693        if ncam == 0 {  // No cameras, ignore.
1694            return;
1695        }
1696
1697        self.camera.fix((self.camera.fixedcamid + direction).rem_euclid(ncam as i32) as usize);
1698    }
1699
1700    /// Toggles full screen mode.
1701    fn toggle_full_screen(&mut self) {
1702        let window = &self.adapter.state.as_ref().unwrap().window;
1703        if window.fullscreen().is_some() {
1704            window.set_fullscreen(None);
1705        }
1706        else {
1707            window.set_fullscreen(Some(Fullscreen::Borderless(None)));
1708        }
1709    }
1710
1711    /// Processes scrolling events.
1712    fn process_scroll(&mut self, change: f64) {
1713        self.camera.move_(
1714            MjtMouse::mjMOUSE_ZOOM,
1715            self.shared_state.lock_unpoison().data_passive.model(),
1716            0.0, -0.05 * change, &self.scene
1717        );
1718    }
1719
1720    /// Processes camera and perturbation movements.
1721    fn process_cursor_pos(&mut self, x: f64, y: f64) {
1722        self.raw_cursor_position = (x, y);
1723        // Calculate the change in mouse position since last call
1724        let dx = x - self.last_x;
1725        let dy = y - self.last_y;
1726        self.last_x = x;
1727        self.last_y = y;
1728        let window = &self.adapter.state.as_ref().unwrap().window;
1729        let modifiers = &self.modifiers.state();
1730        let buttons = &self.buttons_pressed;
1731        let shift = modifiers.shift_key();
1732
1733        // Check mouse presses and move the camera if any of them is pressed
1734        let action;
1735        let height = window.inner_size().height as f64;
1736
1737        let mut lock = self.shared_state.lock_unpoison();
1738        let ViewerSharedState {data_passive, pert, ..} = lock.deref_mut();
1739        if buttons.contains(ButtonsPressed::LEFT) {
1740            if pert.active == MjtPertBit::mjPERT_TRANSLATE as i32 {
1741                action = if shift {MjtMouse::mjMOUSE_MOVE_H} else {MjtMouse::mjMOUSE_MOVE_V};
1742            }
1743            else {
1744                action = if shift {MjtMouse::mjMOUSE_ROTATE_H} else {MjtMouse::mjMOUSE_ROTATE_V};
1745            }
1746        }
1747        else if buttons.contains(ButtonsPressed::RIGHT) {
1748            action = if shift {MjtMouse::mjMOUSE_MOVE_H} else {MjtMouse::mjMOUSE_MOVE_V};
1749        }
1750        else if buttons.contains(ButtonsPressed::MIDDLE) {
1751            action = MjtMouse::mjMOUSE_ZOOM;
1752        }
1753        else {
1754            return;  // If buttons aren't pressed, ignore.
1755        }
1756
1757        // When the perturbation isn't active, move the camera
1758        if pert.active == 0 {
1759            self.camera.move_(
1760                action,
1761                data_passive.model(),
1762                dx / height, dy / height, &self.scene
1763            );
1764        }
1765        else {  // When the perturbation is active, move apply the perturbation.
1766            pert.move_(data_passive, action, dx / height, dy / height, &self.scene);
1767        }
1768    }
1769
1770    /// Processes left clicks and double left clicks.
1771    fn process_left_click(&mut self, state: ElementState) {
1772        let modifier_state = self.modifiers.state();
1773        let mut lock = self.shared_state.lock_unpoison();
1774        let ViewerSharedState {data_passive, pert, ..} = lock.deref_mut();
1775        match state {
1776            ElementState::Pressed => {
1777                // Clicking and holding applies perturbation
1778                if pert.select > 0 && modifier_state.control_key() {
1779                    let type_ = if modifier_state.alt_key() {
1780                        MjtPertBit::mjPERT_TRANSLATE
1781                    } else {
1782                        MjtPertBit::mjPERT_ROTATE
1783                    };
1784                    pert.start(type_, data_passive, &self.scene);
1785                }
1786
1787                // Double click detection
1788                if self.last_bnt_press_time.elapsed().as_millis() < DOUBLE_CLICK_WINDOW_MS {
1789                    let cp = self.raw_cursor_position;
1790                    let x = cp.0;
1791                    let y = self.rect_full.height as f64 - cp.1;
1792
1793                    // Obtain the selection 
1794                    let rect = &self.rect_full;
1795                    let sel = self.scene.find_selection(
1796                        data_passive, &self.opt,
1797                        rect.width as MjtNum / rect.height as MjtNum,
1798                        (x - rect.left as MjtNum) / rect.width as MjtNum,
1799                        (y - rect.bottom as MjtNum) / rect.height as MjtNum
1800                    );
1801
1802                    // Set tracking camera
1803                    if modifier_state.alt_key() {
1804                        if let Some(body_id) = sel.body_id {
1805                            self.camera.lookat = sel.point;
1806                            if modifier_state.control_key() {
1807                                self.camera.track(body_id);
1808                            }
1809                        }
1810                    }
1811                    else {
1812                        // Mark selection
1813                        if let Some(body_id) = sel.body_id {
1814                            pert.select = body_id as i32;
1815                            pert.flexselect = sel.flex_id.map(|v| v as i32).unwrap_or(-1);
1816                            pert.skinselect = sel.skin_id.map(|v| v as i32).unwrap_or(-1);
1817                            pert.active = 0;
1818                            pert.update_local_pos(&sel.point, data_passive);
1819                        }
1820                        else {
1821                            pert.select = -1;
1822                            pert.flexselect = -1;
1823                            pert.skinselect = -1;
1824                        }
1825                    }
1826                }
1827                self.last_bnt_press_time = Instant::now();
1828            },
1829            ElementState::Released => {
1830                // Clear perturbation when left click is released.
1831                pert.active = 0;
1832            },
1833        };
1834    }
1835}
1836
1837/// Releases OpenGL resources (rendering context and egui painter) while the
1838/// GL context is still current.
1839impl Drop for MjViewer {
1840    fn drop(&mut self) {
1841        // Ensure the GL context is current before the implicit field drops so that
1842        // MjrContext::drop (which calls mjr_freeContext) can properly free OpenGL resources.
1843        if let Some(ref state) = self.adapter.state {
1844            let _ = state.gl_context.make_current(&state.gl_surface);
1845        }
1846
1847        // Release egui painter GL resources while the context is still current.
1848        #[cfg(feature = "viewer-ui")]
1849        self.ui.destroy_gl();
1850    }
1851}
1852
1853
1854/// Builder for [`MjViewer`].
1855/// ### Default settings:
1856/// - `window_name`: MuJoCo Rust Viewer (MuJoCo \<MuJoCo version here\>)
1857/// - `max_user_geoms`: 0
1858/// - `vsync`: false
1859/// - `warn_non_realtime`: false
1860/// 
1861#[derive(Debug, Clone)]
1862pub struct MjViewerBuilder {
1863    /// The name shown on the window decoration.
1864    window_name: Cow<'static, str>,
1865    /// Maximum number of geoms that can be given by the user for custom visualization.
1866    max_user_geoms: usize,
1867    /// Start the viewer with vertical synchronization. This should be used only if rendering
1868    /// and simulation are separated by threads and you are ok with [`MjViewer::render`]
1869    /// blocking to achieve the correct refresh rate (of your monitor).
1870    vsync: bool,
1871
1872    /// Start the viewer with warnings enabled for non-realtime synchronization.
1873    /// When this is enabled and the simulation state isn't synced in realtime, an overlay will be displayed
1874    /// in the bottom right corner indicating the realtime percentage.
1875    /// The warning will only be shown if the deviation is 2 % from realtime or more.
1876    warn_non_realtime: bool,
1877}
1878
1879impl MjViewerBuilder {
1880    builder_setters! {
1881        window_name: S where S: Into<Cow<'static, str>>; "text shown in the title of the window.";
1882        max_user_geoms: usize; "maximum number of geoms that can be drawn by the user in addition to the regular geoms.";
1883        vsync: bool; "enable vertical synchronization by default.";
1884        warn_non_realtime: bool; "enable showing an overlay when the simulation state isn't synced in realtime (deviation larger than 2 %).";
1885    }
1886}
1887
1888impl MjViewerBuilder {
1889    /// Creates a [`MjViewerBuilder`] with default settings.
1890    pub fn new() -> Self {
1891        Self {
1892            window_name: Cow::Owned(format!("MuJoCo Rust Viewer (MuJoCo {})", mujoco_version())),
1893            max_user_geoms: 0, vsync: false, warn_non_realtime: false,
1894        }
1895    }
1896
1897    /// Builds a [`MjViewer`] with the configured options.
1898    /// # Returns
1899    /// On success, returns [`Ok`] variant containing the [`MjViewer`].
1900    /// # Errors
1901    /// - [`MjViewerError::EventLoopError`] if the event loop fails to initialize.
1902    /// - [`MjViewerError::GlInitFailed`] if OpenGL / window initialization fails.
1903    /// - [`MjViewerError::GlutinError`] if a glutin operation fails.
1904    /// - [`MjViewerError::PainterInitError`] if the UI painter fails to initialize
1905    ///   (feature `viewer-ui`).
1906    pub fn build_passive<M: Deref<Target = MjModel>>(&self, model: M) -> Result<MjViewer, MjViewerError> {
1907        let (w, h) = MJ_VIEWER_DEFAULT_SIZE_PX;
1908        let mut event_loop = EventLoop::new().map_err(MjViewerError::EventLoopError)?;
1909        let adapter = RenderBase::new(
1910            w, h,
1911            self.window_name.to_string(),
1912            &mut event_loop,
1913            true  // process events
1914        )?;
1915
1916        // Initialize the OpenGL related things
1917        let RenderBaseGlState {
1918            gl_context,
1919            gl_surface,
1920            #[cfg(feature = "viewer-ui")] window,
1921            ..
1922        } = adapter.state.as_ref().unwrap();
1923        gl_context.make_current(gl_surface).map_err(MjViewerError::GlutinError)?;
1924
1925        // Configure vertical synchronization
1926        if self.vsync {
1927            gl_surface.set_swap_interval(
1928                gl_context,
1929                glutin::surface::SwapInterval::Wait(NonZero::<u32>::MIN)
1930            ).map_err(MjViewerError::GlutinError)?;
1931        } else {
1932            gl_surface.set_swap_interval(gl_context, glutin::surface::SwapInterval::DontWait).map_err(
1933                MjViewerError::GlutinError
1934            )?;
1935        }
1936
1937        event_loop.set_control_flow(winit::event_loop::ControlFlow::Poll);
1938
1939        let ngeom = model.ffi().ngeom as usize;
1940        let scene = MjvScene::new(&*model, ngeom + self.max_user_geoms + EXTRA_SCENE_GEOM_SPACE);
1941        // SAFETY: The OpenGL context was made current above via gl_surface.
1942        let context = unsafe { MjrContext::new(&model) };
1943        let camera  = MjvCamera::new_free(&model);
1944
1945        // Tracking of changes made between syncs
1946        let shared_state = Arc::new(Mutex::new(ViewerSharedState::new(&*model, self.max_user_geoms)));
1947        let running_flag = shared_state.lock_unpoison().running.clone();
1948
1949        // User interface
1950        #[cfg(feature = "viewer-ui")]
1951        let ui = ui::ViewerUI::new(&model, window, &gl_surface.display())?;
1952        #[cfg(feature = "viewer-ui")]
1953        let mut status = ViewerStatusBit::UI;
1954        #[cfg(not(feature = "viewer-ui"))]
1955        let mut status = ViewerStatusBit::HELP;
1956
1957        status.set(ViewerStatusBit::VSYNC, self.vsync);
1958        status.set(ViewerStatusBit::WARN_REALTIME, self.warn_non_realtime);
1959
1960        Ok(MjViewer {
1961            scene,
1962            context,
1963            camera,
1964            opt: MjvOption::default(),
1965            ncam: model.ffi().ncam,
1966            shared_state,
1967            running_flag,
1968            last_x: 0.0,
1969            last_y: 0.0,
1970            last_bnt_press_time: Instant::now(),
1971            fps_timer: Instant::now(),
1972            fps_smooth: 60.0,
1973            rect_view: MjrRectangle::default(),
1974            rect_full: MjrRectangle::default(),
1975            adapter,
1976            event_loop,
1977            modifiers: Modifiers::default(),
1978            buttons_pressed: ButtonsPressed::empty(),
1979            raw_cursor_position: (0.0, 0.0),
1980            #[cfg(feature = "viewer-ui")] ui,
1981            status,
1982            screenshot_pending: None
1983        })
1984    }
1985}
1986
1987/// Delegates to [`MjViewerBuilder::new`].
1988impl Default for MjViewerBuilder {
1989    fn default() -> Self {
1990        MjViewerBuilder::new()
1991    }
1992}
1993
1994bitflags! {
1995    /// Internal bit-flags that track the visibility state of various on-screen overlays.
1996    #[derive(Debug)]
1997    struct ViewerStatusBit: u8 {
1998        const HELP = 1 << 0;
1999        const VSYNC = 1 << 1;
2000        const INFO = 1 << 2;
2001        const WARN_REALTIME = 1 << 3;
2002        #[cfg(feature = "viewer-ui")] const UI = 1 << 4;
2003    }
2004}
2005
2006bitflags! {
2007    /// Boolean flags for tracking button press events.
2008    #[derive(Debug)]
2009    struct ButtonsPressed: u8 {
2010        const LEFT = 1 << 0;
2011        const MIDDLE = 1 << 1;
2012        const RIGHT = 1 << 2;
2013    }
2014}