ledstick-rs/src/main.rs
2025-03-19 01:48:48 +01:00

427 lines
14 KiB
Rust

use bytemuck::{bytes_of, bytes_of_mut, Pod, Zeroable};
use esp_idf_svc::{
hal::{
delay::FreeRtos, gpio::AnyIOPin, i2s, peripherals::Peripherals, spi, units::FromValueType,
},
sys::{esp_dsp, esp_nofail, esp_random, TickType_t},
};
use anyhow::{bail, Result};
const LED_COUNT: usize = 72;
const AUDIO_SAMPLES_PER_BUF: usize = 1024;
const AUDIO_BUFFERS: usize = 2;
const AUDIO_BANDS: usize = 3;
type AudioBuffer = [i32; AUDIO_SAMPLES_PER_BUF];
type DspBuffer = [f32; AUDIO_SAMPLES_PER_BUF];
#[derive(Clone, Copy, Eq, PartialEq, Pod, Zeroable)]
#[repr(C, align(4))]
struct Rgbv {
r: u8,
g: u8,
b: u8,
_o: u8,
}
impl Rgbv {
const _O_ONES: u8 = 0xE0;
#[rustfmt::skip] const fn black(o: u8) -> Self { assert!(o<=Self::MAX_O); Self {r: 0x00, g: 0x00, b: 0x00, _o: Self::_O_ONES | o } }
#[rustfmt::skip] const fn white(o: u8) -> Self { assert!(o<=Self::MAX_O); Self {r: 0xFF, g: 0xFF, b: 0xFF, _o: Self::_O_ONES | o } }
#[rustfmt::skip] const fn red(o: u8) -> Self { assert!(o<=Self::MAX_O); Self {r: 0xFF, g: 0x00, b: 0x00, _o: Self::_O_ONES | o } }
#[rustfmt::skip] const fn green(o: u8) -> Self { assert!(o<=Self::MAX_O); Self {r: 0x00, g: 0xFF, b: 0x00, _o: Self::_O_ONES | o } }
#[rustfmt::skip] const fn blue(o: u8) -> Self { assert!(o<=Self::MAX_O); Self {r: 0x00, g: 0x00, b: 0xFF, _o: Self::_O_ONES | o } }
#[rustfmt::skip] const fn cyan(o: u8) -> Self { assert!(o<=Self::MAX_O); Self {r: 0x00, g: 0xFF, b: 0xFF, _o: Self::_O_ONES | o } }
#[rustfmt::skip] const fn orange(o: u8) -> Self { assert!(o<=Self::MAX_O); Self {r: 0xFF, g: 0x80, b: 0x00, _o: Self::_O_ONES | o } }
#[rustfmt::skip] const fn yellow(o: u8) -> Self { assert!(o<=Self::MAX_O); Self {r: 0xFF, g: 0xFF, b: 0x00, _o: Self::_O_ONES | o } }
#[rustfmt::skip] const fn pink(o: u8) -> Self { assert!(o<=Self::MAX_O); Self {r: 0xFF, g: 0x00, b: 0xFF, _o: Self::_O_ONES | o } }
const MAX_O: u8 = 31;
pub fn new(r: u8, g: u8, b: u8, o: u8) -> Self {
assert!(o <= Self::MAX_O);
Self {
r,
g,
b,
_o: o | Self::_O_ONES,
}
}
pub fn o(self) -> u8 {
self._o & !Self::_O_ONES
}
#[inline(always)]
pub fn set_o(mut self, o: u8) -> Self {
assert!(o <= Self::MAX_O);
self._o = o | Self::_O_ONES;
self
}
#[inline(always)]
pub fn increase(mut self, r: u8, g: u8, b: u8, o: u8) -> Self {
self.r = self.r.saturating_add(r);
self.g = self.g.saturating_add(g);
self.b = self.b.saturating_add(b);
self.set_o(std::cmp::min(self.o() + o, Self::MAX_O));
self
}
#[inline(always)]
pub fn decrease(mut self, r: u8, g: u8, b: u8, o: u8) -> Self {
self.r = self.r.saturating_sub(r);
self.g = self.g.saturating_sub(g);
self.b = self.b.saturating_sub(b);
self.set_o(self.o().saturating_sub(o));
self
}
/// Converts hue, saturation, value to RGB
/// // copied from rmt_neopixel example
pub fn from_hsv(h: u32, s: u32, v: u32, o: u8) -> Result<Self> {
assert!(o <= Self::MAX_O);
if h > 360 || s > 100 || v > 100 {
bail!("The given HSV values are not in valid range");
}
let s = s as f64 / 100.0;
let v = v as f64 / 100.0;
let c = s * v;
let x = c * (1.0 - (((h as f64 / 60.0) % 2.0) - 1.0).abs());
let m = v - c;
let (r, g, b) = match h {
0..=59 => (c, x, 0.0),
60..=119 => (x, c, 0.0),
120..=179 => (0.0, c, x),
180..=239 => (0.0, x, c),
240..=299 => (x, 0.0, c),
_ => (c, 0.0, x),
};
Ok(Self {
r: ((r + m) * 255.0) as u8,
g: ((g + m) * 255.0) as u8,
b: ((b + m) * 255.0) as u8,
_o: o | Self::_O_ONES,
})
}
}
type LedColors = [Rgbv; LED_COUNT];
#[repr(C, align(4))]
#[derive(Clone, Copy, Eq, PartialEq, Pod, Zeroable)]
struct LedData {
zeros: u32,
leds: LedColors,
ones: u32,
}
impl LedData {
pub fn new() -> Self {
Self {
zeros: 0,
leds: [Rgbv::new(0, 0, 0, 0); LED_COUNT],
ones: 0, // sic. works as well, and triggers PWM change at the end of frame transfer instead of next one.
}
}
}
fn falloff(old: i32, new: i32) -> i32 {
(old >> 1) + (old >> 2) + (new >> 2)
}
fn falloff_f(old: f32, new: f32) -> f32 {
(old / 2.0f32) + (old / 4.0f32) + (new / 4.0f32)
}
fn random_at_most(max: u32) -> u32 {
// impl from https://stackoverflow.com/a/6852396, adapted to uint32/2
// Assumes 0 <= max <= INT32_MAX
// Returns in the closed interval [0, max]
assert!(max < u32::MAX);
let num_bins = max + 1;
let num_rand = i32::MAX as u32 + 1;
let bin_size = num_rand / num_bins;
let defect = num_rand % num_bins;
let mut x: u32;
loop {
unsafe {
x = esp_random() >> 1; // This is carefully written not to overflow
if num_rand - defect > x {
break;
}
}
}
// Truncated division is intentional
x / bin_size
}
struct AudioProcessor {
floating_max: i32,
current_powers: [f32; AUDIO_BANDS],
avg_powers: [f32; AUDIO_BANDS],
fft_buffer: [DspBuffer; AUDIO_BUFFERS],
next_fft_buf: usize,
fft_window: DspBuffer,
}
impl AudioProcessor {
pub fn new() -> Self {
let mut buf: DspBuffer = [0f32; AUDIO_SAMPLES_PER_BUF];
unsafe {
esp_dsp::dsps_wind_hann_f32(buf.as_mut_ptr(), AUDIO_SAMPLES_PER_BUF as i32);
}
AudioProcessor {
floating_max: 0i32,
current_powers: [0f32; AUDIO_BANDS],
avg_powers: [0f32; AUDIO_BANDS],
fft_buffer: [[0f32; AUDIO_SAMPLES_PER_BUF]; AUDIO_BUFFERS],
next_fft_buf: 0,
fft_window: buf,
}
}
pub fn process(&mut self, audio: &AudioBuffer) -> usize {
let &(mut proc_fft_buffer) = &self.fft_buffer[self.next_fft_buf];
/* calculate floating max */
let mut new_max = 0i32;
for value in audio {
new_max = std::cmp::max(new_max, value.saturating_abs());
}
/* get maximum */
self.floating_max = std::cmp::max(
10000000,
if new_max > self.floating_max {
new_max
} else {
falloff(self.floating_max, falloff(self.floating_max, new_max))
},
);
/* convert to floats for input to fft */
for it in audio.iter().zip(proc_fft_buffer.iter_mut()) {
let (audio_it, fft_it) = it;
*fft_it = (*audio_it as f32) / (i32::MAX as f32);
}
/* do fft */
let half_sample_count = (AUDIO_SAMPLES_PER_BUF / 2) as i32;
unsafe {
esp_nofail!(esp_dsp::dsps_mul_f32_ae32(
proc_fft_buffer.as_ptr(),
self.fft_window.as_ptr(),
proc_fft_buffer.as_mut_ptr(),
AUDIO_SAMPLES_PER_BUF as i32,
1,
1,
1,
));
esp_nofail!(esp_dsp::dsps_fft2r_fc32_aes3_(
proc_fft_buffer.as_mut_ptr(),
half_sample_count,
esp_dsp::dsps_fft_w_table_fc32,
)); // operating on half length but complex
esp_nofail!(esp_dsp::dsps_bit_rev2r_fc32(
proc_fft_buffer.as_mut_ptr(),
half_sample_count
)); // operating on half length but complex
esp_nofail!(esp_dsp::dsps_cplx2real_fc32_ae32_(
proc_fft_buffer.as_mut_ptr(),
half_sample_count,
esp_dsp::dsps_fft_w_table_fc32,
esp_dsp::dsps_fft_w_table_size,
)); // operating on half length but complex
for i in 0..half_sample_count as usize {
proc_fft_buffer[i] = (proc_fft_buffer[i * 2] * proc_fft_buffer[i * 2]
+ proc_fft_buffer[i * 2 + 1] * proc_fft_buffer[i * 2 + 1])
.sqrt();
}
}
/* do band stats */
self.current_powers[0] = proc_fft_buffer[1..8].iter().sum::<f32>() / 8f32;
self.current_powers[1] = proc_fft_buffer[9..86].iter().sum::<f32>() / 78f32;
self.current_powers[2] = proc_fft_buffer[87..470].iter().sum::<f32>() / 384f32;
for it in self.current_powers.iter().zip(self.avg_powers.iter_mut()) {
let (current, avg) = it;
*avg = falloff_f(*avg, *current);
}
let last_fft_buf = self.next_fft_buf;
self.next_fft_buf = (self.next_fft_buf + 1) % AUDIO_BUFFERS;
last_fft_buf
}
}
trait LedEffect {
fn render(&mut self, processed: &AudioProcessor, fft_buf: usize, leds: &LedColors);
}
struct LedEffectBassSparks {}
impl LedEffect for LedEffectBassSparks {
fn render(&mut self, processed: &AudioProcessor, _fft_buf: usize, &(mut leds): &LedColors) {
let bass_color = if processed.floating_max > 10100000
&& (processed.current_powers[0] > 1.25 * processed.avg_powers[0])
{
Rgbv::new(127, 0, 255, 4)
} else {
Rgbv::new(0, 0, 0, 0)
};
leds.fill(bass_color);
bass_color.decrease(3, 5, 5, 0);
if true
/*processed.floating_max > 10100000
&& (processed.current_powers[1] > 1.35 * processed.avg_powers[1])
&& (processed.current_powers[2] > 1.35 * processed.avg_powers[2])*/
{
for _ in 0..10 {
let led_index = random_at_most(LED_COUNT as u32 - 1) as usize;
leds[led_index] = Rgbv::white(31);
}
}
}
}
fn main() -> anyhow::Result<()> {
// It is necessary to call this function once. Otherwise some patches to the runtime
// implemented by esp-idf-sys might not link properly. See https://github.com/esp-rs/esp-idf-template/issues/71
esp_idf_svc::sys::link_patches();
// Bind the log crate to the ESP Logging facilities
esp_idf_svc::log::EspLogger::initialize_default();
let peripherals = Peripherals::take().unwrap();
// leds
let mut leds = LedData::new();
// audio buffers
let mut audio: [AudioBuffer; AUDIO_BUFFERS] = [[0; AUDIO_SAMPLES_PER_BUF]; AUDIO_BUFFERS];
let mut next_audio_buf: usize = 0;
// interfaces
let led_spi_per = peripherals.spi2;
let mic_i2s_per = peripherals.i2s0;
// pins
let led_spi_sdo = peripherals.pins.gpio11;
let led_spi_sck = peripherals.pins.gpio12;
let mic_i2s_sd = peripherals.pins.gpio5;
let mic_i2s_sclk = peripherals.pins.gpio4;
let mic_i2s_ws = peripherals.pins.gpio6;
// i2s config
let mic_i2s_std_cfg = i2s::config::StdConfig::new(
i2s::config::Config::new().role(i2s::config::Role::Controller),
i2s::config::StdClkConfig::new(
24000,
i2s::config::ClockSource::default(),
i2s::config::MclkMultiple::M256,
),
i2s::config::StdSlotConfig::msb_slot_default(
i2s::config::DataBitWidth::Bits32,
i2s::config::SlotMode::Stereo,
)
.data_bit_width(i2s::config::DataBitWidth::Bits32)
.slot_bit_width(i2s::config::SlotBitWidth::Bits32)
.slot_mode_mask(i2s::config::SlotMode::Stereo, i2s::config::StdSlotMask::Both)
.ws_width(32)
.ws_polarity(false)
.bit_shift(true)
.left_align(true)
.big_endian(false)
.bit_order_lsb(false),
i2s::config::StdGpioConfig::new(false, false, false),
);
let mut mic_drv = i2s::I2sDriver::new_std_rx(
mic_i2s_per,
&mic_i2s_std_cfg,
mic_i2s_sclk,
mic_i2s_sd,
AnyIOPin::none(),
mic_i2s_ws,
)?;
// spi config
let mut led_drv = spi::SpiDeviceDriver::new_single(
led_spi_per,
led_spi_sck,
led_spi_sdo,
AnyIOPin::none(),
AnyIOPin::none(),
&spi::config::DriverConfig::new(),
&spi::config::Config::new()
.baudrate(1.MHz().into())
.data_mode(spi::config::MODE_3),
)?;
unsafe {
esp_nofail!(esp_dsp::dsps_fft2r_init_fc32(
std::ptr::null_mut(),
(AUDIO_SAMPLES_PER_BUF / 2) as i32
));
esp_nofail!(esp_dsp::dsps_fft4r_init_fc32(
std::ptr::null_mut(),
(AUDIO_SAMPLES_PER_BUF / 2) as i32
));
}
let mut processor = AudioProcessor::new();
let mut effect = LedEffectBassSparks {};
// loop {
// leds.leds[0] = Rgbv::red(4);
// let output_buffer = bytes_of(&leds);
// led_drv.write(output_buffer)?;
// FreeRtos::delay_ms(10);
// }
mic_drv.rx_enable()?;
loop {
// let buffer: &mut [u8; AUDIO_SAMPLES_PER_BUF*4] = cast_slice_mut(&mut audio[next_audio_buf]);
// let buffer = bytes_of_mut(&mut audio[next_audio_buf]);
let mut buffer:[u8;AUDIO_SAMPLES_PER_BUF*4] = [0;AUDIO_SAMPLES_PER_BUF*4];
let num_bytes_read = mic_drv.read(buffer.as_mut_slice(), TickType_t::MAX)?;
if num_bytes_read != AUDIO_SAMPLES_PER_BUF * 4 {
log::error!("buffer underflow");
}
for i in 0..AUDIO_SAMPLES_PER_BUF {
let sample:&[u8;4] = &buffer[i*4..i*4+4].try_into().expect("bla");
audio[next_audio_buf][i] = i32::from_be_bytes(*sample);
}
// log::info!("a: {:08x}", audio[next_audio_buf][0]);
let current_fft_buf = processor.process(&audio[next_audio_buf]);
effect.render(&processor, current_fft_buf, &leds.leds);
let output_buffer = bytes_of(&leds);
led_drv.write(output_buffer)?;
next_audio_buf = (next_audio_buf + 1) % AUDIO_BUFFERS;
FreeRtos::delay_ms(10);
}
}