#![forbid(unsafe_code)]
#![warn(missing_docs)]
extern crate rustfft;
extern crate byteorder;
extern crate rubato;
use byteorder::{LittleEndian, ReadBytesExt};
use rubato::Resampler;
use rustfft::{num_complex::Complex, num_traits::Zero, Fft, FftPlanner};
use std::fmt::{Debug, Formatter};
use std::ops::{Add, Sub};
use std::path::PathBuf;
use std::sync::Arc;
use std::{
fs::File,
io::{BufReader, Error, Read},
path::Path,
};
#[derive(Copy, Clone)]
pub struct Vec3 {
pub x: f32,
pub y: f32,
pub z: f32,
}
impl Vec3 {
pub fn new(x: f32, y: f32, z: f32) -> Self {
Self { x, y, z }
}
fn dot(self, other: Self) -> f32 {
self.x * other.x + self.y * other.y + self.z * other.z
}
fn cross(self, other: Self) -> Vec3 {
Self {
x: self.y * other.z - self.z * other.y,
y: self.z * other.x - self.x * other.z,
z: self.x * other.y - self.y * other.x,
}
}
fn normalize(self) -> Vec3 {
let i = 1.0 / self.dot(self).sqrt();
Vec3 {
x: self.x * i,
y: self.y * i,
z: self.z * i,
}
}
fn scale(self, k: f32) -> Self {
Self {
x: self.x * k,
y: self.y * k,
z: self.z * k,
}
}
fn lerp(self, other: Self, t: f32) -> Self {
Self {
x: lerpf(self.x, other.x, t),
y: lerpf(self.y, other.y, t),
z: lerpf(self.z, other.z, t),
}
}
}
impl Add for Vec3 {
type Output = Self;
fn add(self, rhs: Self) -> Self::Output {
Self {
x: self.x + rhs.x,
y: self.y + rhs.y,
z: self.z + rhs.z,
}
}
}
impl Sub for Vec3 {
type Output = Self;
fn sub(self, rhs: Self) -> Self::Output {
Self {
x: self.x - rhs.x,
y: self.y - rhs.y,
z: self.z - rhs.z,
}
}
}
fn lerpf(a: f32, b: f32, t: f32) -> f32 {
a + (b - a) * t
}
struct BaryCoords {
u: f32,
v: f32,
w: f32,
}
impl BaryCoords {
fn inside(&self) -> bool {
(self.u >= 0.0) && (self.v >= 0.0) && (self.u + self.v < 1.0)
}
}
fn get_barycentric_coords(p: Vec3, a: Vec3, b: Vec3, c: Vec3) -> BaryCoords {
let v0 = b - a;
let v1 = c - a;
let v2 = p - a;
let d00 = v0.dot(v0);
let d01 = v0.dot(v1);
let d11 = v1.dot(v1);
let d20 = v2.dot(v0);
let d21 = v2.dot(v1);
let denom = d00 * d11 - d01 * d01;
let v = (d11 * d20 - d01 * d21) / denom;
let w = (d00 * d21 - d01 * d20) / denom;
let u = 1.0 - v - w;
BaryCoords { u, v, w }
}
fn ray_triangle_intersection(origin: Vec3, dir: Vec3, vertices: &[Vec3; 3]) -> Option<BaryCoords> {
let ba = vertices[1] - vertices[0];
let ca = vertices[2] - vertices[0];
let normal = ba.cross(ca).normalize();
let d = -vertices[0].dot(normal);
let u = -(origin.dot(normal) + d);
let v = dir.dot(normal);
let t = u / v;
if t >= 0.0 && t <= 1.0 {
let point = origin + dir.scale(t);
let bary = get_barycentric_coords(point, vertices[0], vertices[1], vertices[2]);
if bary.inside() {
return Some(bary);
}
}
None
}
#[derive(Debug)]
pub enum HrtfError {
IoError(std::io::Error),
InvalidFileFormat,
InvalidLength(usize),
}
impl From<std::io::Error> for HrtfError {
fn from(io_err: Error) -> Self {
HrtfError::IoError(io_err)
}
}
#[derive(Clone)]
struct Face {
a: usize,
b: usize,
c: usize,
}
#[derive(Clone)]
pub struct HrtfSphere {
length: usize,
points: Vec<HrtfPoint>,
faces: Vec<Face>,
source: PathBuf,
}
fn make_hrtf(
hrir: Vec<f32>,
pad_length: usize,
planner: &mut FftPlanner<f32>,
) -> Vec<Complex<f32>> {
let mut hrir = hrir
.into_iter()
.map(|s| Complex::new(s, 0.0))
.collect::<Vec<Complex<f32>>>();
for _ in hrir.len()..pad_length {
hrir.push(Complex::zero());
}
planner.plan_fft_forward(pad_length).process(hrir.as_mut());
hrir
}
fn read_hrir(reader: &mut dyn Read, len: usize) -> Result<Vec<f32>, HrtfError> {
let mut hrir = Vec::with_capacity(len);
for _ in 0..len {
hrir.push(reader.read_f32::<LittleEndian>()?);
}
Ok(hrir)
}
fn resample_hrir(hrir: Vec<f32>, ratio: f64) -> Vec<f32> {
if ratio.eq(&1.0) {
hrir
} else {
let params = rubato::InterpolationParameters {
sinc_len: 256,
f_cutoff: 0.95,
oversampling_factor: 160,
interpolation: rubato::InterpolationType::Cubic,
window: rubato::WindowFunction::BlackmanHarris2,
};
let mut resampler = rubato::SincFixedIn::<f32>::new(ratio, params, hrir.len(), 1);
let result = resampler.process(&[hrir]).unwrap();
result.into_iter().next().unwrap()
}
}
fn read_faces(reader: &mut dyn Read, index_count: usize) -> Result<Vec<Face>, HrtfError> {
let mut indices = Vec::with_capacity(index_count);
for _ in 0..index_count {
indices.push(reader.read_u32::<LittleEndian>()?);
}
let faces = indices
.chunks(3)
.map(|f| Face {
a: f[0] as usize,
b: f[1] as usize,
c: f[2] as usize,
})
.collect();
Ok(faces)
}
#[derive(Clone)]
pub struct HrirPoint {
pub pos: Vec3,
left_hrir: Vec<f32>,
right_hrir: Vec<f32>,
}
impl HrirPoint {
pub fn left_hrir(&self) -> &[f32] {
&self.left_hrir
}
pub fn right_hrir(&self) -> &[f32] {
&self.right_hrir
}
}
#[derive(Clone)]
pub struct HrirSphere {
length: usize,
points: Vec<HrirPoint>,
faces: Vec<Face>,
source: PathBuf,
}
impl HrirSphere {
pub fn from_file<P: AsRef<Path>>(path: P, device_sample_rate: u32) -> Result<Self, HrtfError> {
let mut sphere = Self::new(
BufReader::new(File::open(path.as_ref())?),
device_sample_rate,
)?;
sphere.source = path.as_ref().to_owned();
Ok(sphere)
}
pub fn new<R: Read>(mut reader: R, device_sample_rate: u32) -> Result<Self, HrtfError> {
let mut magic = [0; 4];
reader.read_exact(&mut magic)?;
if magic[0] != b'H' && magic[1] != b'R' && magic[2] != b'I' && magic[3] != b'R' {
return Err(HrtfError::InvalidFileFormat);
}
let sample_rate = reader.read_u32::<LittleEndian>()?;
let length = reader.read_u32::<LittleEndian>()? as usize;
if length == 0 {
return Err(HrtfError::InvalidLength(length));
}
let vertex_count = reader.read_u32::<LittleEndian>()? as usize;
let index_count = reader.read_u32::<LittleEndian>()? as usize;
let faces = read_faces(&mut reader, index_count)?;
let ratio = sample_rate as f64 / device_sample_rate as f64;
let mut points = Vec::with_capacity(vertex_count);
for _ in 0..vertex_count {
let x = reader.read_f32::<LittleEndian>()?;
let y = reader.read_f32::<LittleEndian>()?;
let z = reader.read_f32::<LittleEndian>()?;
let left_hrir = resample_hrir(read_hrir(&mut reader, length)?, ratio);
let right_hrir = resample_hrir(read_hrir(&mut reader, length)?, ratio);
points.push(HrirPoint {
pos: Vec3 { x, y, z },
left_hrir,
right_hrir,
});
}
Ok(Self {
points,
length,
faces,
source: Default::default(),
})
}
pub fn transform(&mut self, matrix: &[f32; 16]) {
for pt in self.points.iter_mut() {
let x = pt.pos.x * matrix[0] + pt.pos.y * matrix[4] + pt.pos.z * matrix[8] + matrix[12];
let y = pt.pos.x * matrix[1] + pt.pos.y * matrix[5] + pt.pos.z * matrix[9] + matrix[13];
let z =
pt.pos.x * matrix[2] + pt.pos.y * matrix[6] + pt.pos.z * matrix[10] + matrix[14];
pt.pos.x = x;
pt.pos.y = y;
pt.pos.z = z;
}
}
pub fn points(&self) -> &[HrirPoint] {
&self.points
}
pub fn points_mut(&mut self) -> &mut [HrirPoint] {
&mut self.points
}
pub fn len(&self) -> usize {
self.length
}
}
#[derive(Clone)]
struct HrtfPoint {
pos: Vec3,
left_hrtf: Vec<Complex<f32>>,
right_hrtf: Vec<Complex<f32>>,
}
impl HrtfSphere {
fn new(hrir_sphere: HrirSphere, block_len: usize) -> Self {
let mut planner = FftPlanner::new();
let pad_length = get_pad_len(hrir_sphere.length, block_len);
let points = hrir_sphere
.points
.into_iter()
.map(|p| {
let left_hrtf = make_hrtf(p.left_hrir, pad_length, &mut planner);
let right_hrtf = make_hrtf(p.right_hrir, pad_length, &mut planner);
HrtfPoint {
pos: p.pos,
left_hrtf,
right_hrtf,
}
})
.collect();
Self {
points,
length: hrir_sphere.length,
faces: hrir_sphere.faces,
source: hrir_sphere.source,
}
}
pub fn source(&self) -> &Path {
&self.source
}
fn sample_bilinear(
&self,
left_hrtf: &mut Vec<Complex<f32>>,
right_hrtf: &mut Vec<Complex<f32>>,
dir: Vec3,
) {
let dir = dir.scale(10.0);
for face in self.faces.iter() {
let a = self.points.get(face.a).unwrap();
let b = self.points.get(face.b).unwrap();
let c = self.points.get(face.c).unwrap();
if let Some(bary) = ray_triangle_intersection(
Vec3 {
x: 0.0,
y: 0.0,
z: 0.0,
},
dir,
&[a.pos, b.pos, c.pos],
) {
let len = a.left_hrtf.len();
left_hrtf.clear();
for i in 0..len {
left_hrtf.push(
a.left_hrtf[i] * bary.u + b.left_hrtf[i] * bary.v + c.left_hrtf[i] * bary.w,
);
}
right_hrtf.clear();
for i in 0..len {
right_hrtf.push(
a.right_hrtf[i] * bary.u
+ b.right_hrtf[i] * bary.v
+ c.right_hrtf[i] * bary.w,
);
}
}
}
}
}
#[inline]
fn copy_replace(prev_samples: &mut Vec<f32>, raw_buffer: &mut [Complex<f32>], segment_len: usize) {
if prev_samples.len() != segment_len {
*prev_samples = vec![0.0; segment_len];
}
for (prev_sample, raw_sample) in prev_samples.iter().zip(&mut raw_buffer[..segment_len]) {
*raw_sample = Complex::new(*prev_sample, 0.0);
}
let last_start = raw_buffer.len() - segment_len;
for (prev_sample, raw_sample) in prev_samples.iter_mut().zip(&mut raw_buffer[last_start..]) {
*prev_sample = raw_sample.re;
}
}
#[inline]
fn convolve_overlap_save(
in_buffer: &mut [Complex<f32>],
scratch_buffer: &mut [Complex<f32>],
hrtf: &[Complex<f32>],
hrtf_len: usize,
prev_samples: &mut Vec<f32>,
fft: &dyn Fft<f32>,
ifft: &dyn Fft<f32>,
) {
assert_eq!(hrtf.len(), in_buffer.len());
copy_replace(prev_samples, in_buffer, hrtf_len);
fft.process_with_scratch(in_buffer, scratch_buffer);
for (s, h) in in_buffer.iter_mut().zip(hrtf.iter()) {
*s *= *h;
}
ifft.process_with_scratch(in_buffer, scratch_buffer);
}
#[inline]
fn get_pad_len(hrtf_len: usize, block_len: usize) -> usize {
block_len + hrtf_len - 1
}
pub struct HrtfProcessor {
hrtf_sphere: HrtfSphere,
left_in_buffer: Vec<Complex<f32>>,
right_in_buffer: Vec<Complex<f32>>,
scratch_buffer: Vec<Complex<f32>>,
fft: Arc<dyn Fft<f32>>,
ifft: Arc<dyn Fft<f32>>,
left_hrtf: Vec<Complex<f32>>,
right_hrtf: Vec<Complex<f32>>,
block_len: usize,
interpolation_steps: usize,
}
impl Debug for HrtfProcessor {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
writeln!(f, "HrtfProcessor")
}
}
impl Clone for HrtfProcessor {
fn clone(&self) -> Self {
Self {
hrtf_sphere: self.hrtf_sphere.clone(),
left_in_buffer: self.left_in_buffer.clone(),
right_in_buffer: self.right_in_buffer.clone(),
scratch_buffer: self.scratch_buffer.clone(),
fft: self.fft.clone(),
ifft: self.ifft.clone(),
left_hrtf: self.left_hrtf.clone(),
right_hrtf: self.right_hrtf.clone(),
block_len: self.block_len,
interpolation_steps: self.interpolation_steps,
}
}
}
pub trait InterleavedSamples {
fn left(&self) -> f32;
}
impl InterleavedSamples for f32 {
fn left(&self) -> f32 {
*self
}
}
impl InterleavedSamples for (f32, f32) {
fn left(&self) -> f32 {
self.0
}
}
#[inline]
fn get_raw_samples<T: InterleavedSamples>(
source: &[T],
left: &mut [Complex<f32>],
right: &mut [Complex<f32>],
offset: usize,
) {
assert_eq!(left.len(), right.len());
for ((left, right), samples) in left.iter_mut().zip(right.iter_mut()).zip(&source[offset..]) {
let sample = Complex::new(samples.left(), 0.0);
*left = sample;
*right = sample;
}
}
pub struct HrtfContext<'a, 'b, 'c, T: InterleavedSamples> {
pub source: &'a [T],
pub output: &'b mut [(f32, f32)],
pub new_sample_vector: Vec3,
pub prev_sample_vector: Vec3,
pub prev_left_samples: &'c mut Vec<f32>,
pub prev_right_samples: &'c mut Vec<f32>,
pub new_distance_gain: f32,
pub prev_distance_gain: f32,
}
impl HrtfProcessor {
pub fn new(hrir_sphere: HrirSphere, interpolation_steps: usize, block_len: usize) -> Self {
let hrtf_sphere = HrtfSphere::new(hrir_sphere, block_len);
let pad_length = get_pad_len(hrtf_sphere.length, block_len);
let pt = hrtf_sphere.points.first().unwrap();
let left_hrtf = pt.left_hrtf.clone();
let right_hrtf = pt.right_hrtf.clone();
let mut planner = FftPlanner::new();
Self {
hrtf_sphere,
left_in_buffer: vec![Complex::zero(); pad_length],
right_in_buffer: vec![Complex::zero(); pad_length],
scratch_buffer: vec![Complex::zero(); pad_length],
fft: planner.plan_fft_forward(pad_length),
ifft: planner.plan_fft_inverse(pad_length),
left_hrtf,
right_hrtf,
block_len,
interpolation_steps,
}
}
pub fn hrtf_sphere(&self) -> &HrtfSphere {
&self.hrtf_sphere
}
pub fn process_samples<T: InterleavedSamples>(&mut self, context: HrtfContext<T>) {
let HrtfContext {
source,
output,
new_sample_vector: sample_vector,
prev_sample_vector,
prev_left_samples,
prev_right_samples,
prev_distance_gain,
new_distance_gain,
} = context;
let expected_len = self.interpolation_steps * self.block_len;
assert_eq!(expected_len, source.len());
assert!(output.len() >= expected_len);
let new_sampling_vector = sample_vector;
let prev_sampling_vector = prev_sample_vector;
let pad_length = get_pad_len(self.hrtf_sphere.length, self.block_len);
for step in 0..self.interpolation_steps {
let next = step + 1;
let out = &mut output[(step * self.block_len)..(next * self.block_len)];
let t = next as f32 / self.interpolation_steps as f32;
let sampling_vector = prev_sampling_vector.lerp(new_sampling_vector, t);
self.hrtf_sphere.sample_bilinear(
&mut self.left_hrtf,
&mut self.right_hrtf,
sampling_vector,
);
let hrtf_len = self.hrtf_sphere.length - 1;
get_raw_samples(
source,
&mut self.left_in_buffer[hrtf_len..],
&mut self.right_in_buffer[hrtf_len..],
step * self.block_len,
);
convolve_overlap_save(
&mut self.left_in_buffer,
&mut self.scratch_buffer,
&self.left_hrtf,
hrtf_len,
prev_left_samples,
&*self.fft,
&*self.ifft,
);
convolve_overlap_save(
&mut self.right_in_buffer,
&mut self.scratch_buffer,
&self.right_hrtf,
hrtf_len,
prev_right_samples,
&*self.fft,
&*self.ifft,
);
let distance_gain = lerpf(prev_distance_gain, new_distance_gain, t);
let k = distance_gain / (pad_length as f32);
let left_payload = &self.left_in_buffer[hrtf_len..];
let right_payload = &self.right_in_buffer[hrtf_len..];
for ((out_left, out_right), (processed_left, processed_right)) in
out.iter_mut().zip(left_payload.iter().zip(right_payload))
{
*out_left += processed_left.re * k;
*out_right += processed_right.re * k;
}
}
}
}