From 119c7026b477836e64c766d1ae2ebcd2a6086b7c Mon Sep 17 00:00:00 2001
From: Schmarni <marnistromer@gmail.com>
Date: Wed, 16 Jul 2025 21:43:37 +0200
Subject: [PATCH] fix: offset tracking at startup to center properly

Signed-off-by: Schmarni <marnistromer@gmail.com>
---
 src/main.rs            |  4 ++-
 src/tracking_offset.rs | 82 ++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 85 insertions(+), 1 deletion(-)
 create mode 100644 src/tracking_offset.rs

diff --git a/src/main.rs b/src/main.rs
index ca449b9..9b887c9 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -5,6 +5,7 @@ mod core;
 mod nodes;
 mod objects;
 mod session;
+pub mod tracking_offset;
 #[cfg(feature = "wayland")]
 mod wayland;
 
@@ -70,6 +71,7 @@ use tokio::task::JoinError;
 use tracing::metadata::LevelFilter;
 use tracing::{error, info};
 use tracing_subscriber::{EnvFilter, fmt, prelude::*};
+use tracking_offset::TrackingOffsetPlugin;
 use wayland::{Wayland, WaylandPlugin};
 use zbus::Connection;
 use zbus::fdo::ObjectManager;
@@ -398,7 +400,7 @@ fn bevy_loop(
 	// object plugins
 	app.add_plugins((PlaySpacePlugin, HandPlugin, ControllerPlugin, HmdPlugin));
 	// feature plugins
-	app.add_plugins(WaylandPlugin);
+	app.add_plugins((WaylandPlugin, TrackingOffsetPlugin));
 	app.add_systems(PostStartup, move || {
 		ready_notifier.notify_waiters();
 	});
diff --git a/src/tracking_offset.rs b/src/tracking_offset.rs
new file mode 100644
index 0000000..cf2743d
--- /dev/null
+++ b/src/tracking_offset.rs
@@ -0,0 +1,82 @@
+use bevy::{
+	app::{Plugin, Update},
+	ecs::{
+		resource::Resource,
+		schedule::{Condition, IntoScheduleConfigs, common_conditions::resource_exists},
+		system::{Commands, Res, ResMut},
+	},
+	transform::components::Transform,
+};
+use bevy_mod_openxr::{
+	helper_traits::ToTransform as _,
+	poll_events::{OxrEventHandlerExt, OxrEventIn},
+	resources::OxrFrameState,
+	session::OxrSession,
+};
+use bevy_mod_xr::{session::XrSessionCreated, spaces::XrPrimaryReferenceSpace};
+use glam::{Quat, Vec3};
+use openxr::ReferenceSpaceType;
+
+pub struct TrackingRootOffsetPlugin;
+
+impl Plugin for TrackingRootOffsetPlugin {
+	fn build(&self, app: &mut bevy::app::App) {
+		app.add_oxr_event_handler(reset_offset);
+		app.add_systems(XrSessionCreated, |mut cmds: Commands| {
+			cmds.insert_resource(OffsetTag);
+		});
+		app.add_systems(
+			Update,
+			offset.run_if(resource_exists::<OffsetTag>.and(resource_exists::<OxrFrameState>)),
+		);
+	}
+}
+
+#[derive(Resource)]
+struct OffsetTag;
+
+fn reset_offset(
+	oxr_event: OxrEventIn,
+	mut ref_space: ResMut<XrPrimaryReferenceSpace>,
+	session: Res<OxrSession>,
+) {
+	if let openxr::Event::ReferenceSpaceChangePending(v) = *oxr_event
+		&& v.reference_space_type() == ReferenceSpaceType::LOCAL
+	{
+		let space = session
+			.create_reference_space(ReferenceSpaceType::LOCAL, Transform::IDENTITY)
+			.unwrap();
+		session.destroy_space(ref_space.0.0).unwrap();
+		ref_space.0 = space;
+	}
+}
+
+fn offset(
+	session: Res<OxrSession>,
+	state: Res<OxrFrameState>,
+	mut primary_ref_space: ResMut<XrPrimaryReferenceSpace>,
+	mut cmds: Commands,
+) {
+	cmds.remove_resource::<OffsetTag>();
+	let local = session
+		.create_reference_space(ReferenceSpaceType::LOCAL, Transform::IDENTITY)
+		.unwrap();
+	let view = session
+		.create_reference_space(ReferenceSpaceType::VIEW, Transform::IDENTITY)
+		.unwrap();
+	let view_pose = session
+		.locate_space(&view, &local, state.predicted_display_time)
+		.unwrap()
+		.pose
+		.to_transform();
+	let offset = view_pose.with_rotation(Quat::from_axis_angle(
+		Vec3::Y,
+		view_pose.rotation.to_euler(glam::EulerRot::XYZ).1,
+	));
+	let offset = Transform::from_matrix(offset.compute_matrix());
+	let local_offset = session
+		.create_reference_space(ReferenceSpaceType::LOCAL, offset)
+		.unwrap();
+	session.destroy_space(primary_ref_space.0.0).unwrap();
+	primary_ref_space.0 = local_offset;
+}