Drive the nose too

Author: BenChung

deploy/src/main.jl

@@ -6,6 +6,7 @@ include(joinpath(@__DIR__, "mpu6000.jl"))
 include(joinpath(@__DIR__, "encoder.jl"))
 include(joinpath(@__DIR__, "shift_driver.jl"))
 include(joinpath(@__DIR__, "seven_seg.jl"))
+include(joinpath(@__DIR__, "ws2812_driver.jl"))
 
 using .GPIO
 using .PWM
@@ -29,7 +30,9 @@ const PWM_MAX_VALUE = 1024  # PWM duty cycle range (0-1024)
 # Global handles for GPIO pins and PWM channels
 mutable struct HardwareContext
     gpio::GPIO.GPIOController
+    pio::PIOBlock
     chain::ShiftRegisterChain
+    nose::WS2812
     pins::Vector{GPIO.GPIOPin}
 end
 
@@ -42,7 +45,10 @@ Called on Ctrl-C or program exit.
 function shutdown!(hw::HardwareContext)
     println(Core.stdout, "Shutting down hardware...")
     hw.chain[0:23] = false
+    set_color!(hw.nose, 0, 0, 0)
     close(hw.chain)
+    close(hw.nose)
+    close(hw.pio)
     for pin in hw.pins
         close(pin)
     end
@@ -61,6 +67,7 @@ end
 const CM_PRESENT=23
 const D1=5
 const D2=24
+const NOSE_RGB=18
 
 function (@main)(args)::Cint
     println(Core.stdout, "Balance car starting...")
@@ -74,12 +81,20 @@ function (@main)(args)::Cint
     d2 = GPIO.request_output(gpio, D2, "d2", 0)
     println(Core.stdout, "GPIO configured")
 
+    # Open shared PIO block
+    pio = open_pio(0)
+
     # Initialize 3 chained shift registers via PIO
-    chain = open_shift_registers()
+    chain = open_shift_registers(pio)
     chain[0:23] = false
     println(Core.stdout, "Shift registers initialized ($NUM_REGISTERS x 8-bit, $NBITS outputs)")
 
-    hw = HardwareContext(gpio, chain, [cm_present, d1, d2])
+    # Initialize nose RGB LED (WS2812) on same PIO block
+    nose = open_ws2812(pio, NOSE_RGB)
+    set_color!(nose, 0, 0, 0)
+    println(Core.stdout, "Nose RGB LED initialized on pin $NOSE_RGB")
+
+    hw = HardwareContext(gpio, pio, chain, nose, [cm_present, d1, d2])
 
     # Enable hardware
     GPIO.set_value(cm_present, 1)
@@ -93,6 +108,7 @@ function (@main)(args)::Cint
     disp2 = SevenSeg(chain, 16)  # 3rd register: bits 16–23
 
     digit = 0
+    hue = 0
 
     # Control loop timing (500ms = 2Hz)
     loop_period_ns = 500_000_000
@@ -119,6 +135,26 @@ function (@main)(args)::Cint
                 show_digit!(disp1, (digit + 1) % 10)
             end
             digit = (digit + 1) % 10
+
+            # Cycle nose LED through hue wheel
+            # Simple HSV→RGB with S=V=255, varying H
+            h = hue ÷ 43
+            f = (hue - h * 43) * 6
+            q = 255 - f
+            if h == 0
+                set_color!(nose, 255, f, 0)
+            elseif h == 1
+                set_color!(nose, q, 255, 0)
+            elseif h == 2
+                set_color!(nose, 0, 255, f)
+            elseif h == 3
+                set_color!(nose, 0, q, 255)
+            elseif h == 4
+                set_color!(nose, f, 0, 255)
+            else
+                set_color!(nose, 255, 0, q)
+            end
+            hue = (hue + 25) % 256
         end
     catch e
         if e isa InterruptException

deploy/src/shift_driver.jl

@@ -84,42 +84,36 @@ mutable struct ShiftRegisterChain
 end
 
 """
-    open_shift_registers(pio_idx=0) -> ShiftRegisterChain
+    open_shift_registers(pio::PIOBlock) -> ShiftRegisterChain
 
-Initialize the PIO and state machine for driving the shift register chain.
-Caller is responsible for calling `close` when done.
+Initialize a state machine on `pio` for driving the shift register chain.
+The PIO block must already be open (caller manages its lifetime).
 """
-function open_shift_registers(pio_idx::Integer=0)
-    pio = open_pio(pio_idx)
-    try
-        prog, config = shift_register_program(pio;
-            ser_pin=SER_PIN, clk_pin=CLK_PIN, rclk_pin=RCLK_PIN,
-            nbits=NBITS, clkdiv=CLKDIV,
-        )
+function open_shift_registers(pio::PIOBlock)
+    prog, config = shift_register_program(pio;
+        ser_pin=SER_PIN, clk_pin=CLK_PIN, rclk_pin=RCLK_PIN,
+        nbits=NBITS, clkdiv=CLKDIV,
+    )
 
-        for pin in (SER_PIN, CLK_PIN, RCLK_PIN)
-            pio_pin_init!(pio, pin)
-        end
+    for pin in (SER_PIN, CLK_PIN, RCLK_PIN)
+        pio_pin_init!(pio, pin)
+    end
 
-        offset = load_program!(pio, prog, config)
-
-        sm = claim_sm(pio)
-        try
-            pin_mask = UInt32(1) << SER_PIN | UInt32(1) << CLK_PIN | UInt32(1) << RCLK_PIN
-            set_pindirs!(sm, pin_mask, pin_mask)
-            PIOLib.init!(sm, offset, config)
-            setup_shift_register!(sm, NBITS)
-            enable!(sm)
-        catch
-            unclaim!(sm)
-            rethrow()
-        end
+    offset = load_program!(pio, prog, config)
 
-        return ShiftRegisterChain(pio, sm, UInt32(0), false)
+    sm = claim_sm(pio)
+    try
+        pin_mask = UInt32(1) << SER_PIN | UInt32(1) << CLK_PIN | UInt32(1) << RCLK_PIN
+        set_pindirs!(sm, pin_mask, pin_mask)
+        PIOLib.init!(sm, offset, config)
+        setup_shift_register!(sm, NBITS)
+        enable!(sm)
     catch
-        close(pio)
+        unclaim!(sm)
         rethrow()
     end
+
+    return ShiftRegisterChain(pio, sm, UInt32(0), false)
 end
 
 # --- Flush ---
@@ -265,5 +259,4 @@ end
 function Base.close(chain::ShiftRegisterChain)
     disable!(chain.sm)
     unclaim!(chain.sm)
-    close(chain.pio)
 end

deploy/src/ws2812_driver.jl

@@ -0,0 +1,75 @@
+using PIOLib
+
+function ws2812_program(pio::PIOBlock; pin::Integer, sys_clk::Real=200_000_000.0)
+    T1 = 2
+    T2 = 5
+    T3 = 3
+    cycles_per_bit = T1 + T2 + T3
+    div = clkdiv(800_000.0 * cycles_per_bit; sys_clk)
+
+    prog = build_program([
+        WrapTarget(),
+        Label(:bitloop),
+        Out(RegX(), 1; sideset=0, delay=T3-1),
+        Jmp{:not_x}(:do_zero; sideset=1, delay=T1-1),
+        Jmp{:always}(:bitloop; sideset=1, delay=T2-1),
+        Label(:do_zero),
+        Nop(; sideset=0, delay=T2-1),
+        Wrap(),
+    ]; sideset_bits=1)
+
+    config = SMConfig(pio;
+        sideset_pin_base=pin,
+        sideset=(1, false, false),
+        out_shift=(false, true, 24),
+        clkdiv=div,
+        wrap=(prog.wrap_target, prog.wrap),
+        fifo_join=PIOLib.LibPIO.PIO_FIFO_JOIN_TX,
+    )
+
+    prog, config
+end
+
+ws2812_put!(sm::StateMachine, grb::UInt32) = put!(sm, grb << 8)
+
+function ws2812_rgb!(sm::StateMachine, r::Integer, g::Integer, b::Integer)
+    grb = (UInt32(g) << 16) | (UInt32(r) << 8) | UInt32(b)
+    ws2812_put!(sm, grb)
+end
+
+struct WS2812
+    pio::PIOBlock
+    sm::StateMachine
+end
+
+"""
+    open_ws2812(pio::PIOBlock, pin::Integer) -> WS2812
+
+Initialize a WS2812 LED on `pin` using an SM from the given PIO block.
+The PIO block must already be open (shared with other drivers).
+"""
+function open_ws2812(pio::PIOBlock, pin::Integer)
+    prog, config = ws2812_program(pio; pin=pin)
+
+    pio_pin_init!(pio, pin)
+    offset = load_program!(pio, prog, config)
+
+    sm = claim_sm(pio)
+    try
+        set_consecutive_pindirs!(sm, pin, 1, true)
+        PIOLib.init!(sm, offset, config)
+        enable!(sm)
+    catch
+        unclaim!(sm)
+        rethrow()
+    end
+
+    WS2812(pio, sm)
+end
+
+set_color!(led::WS2812, r::Integer, g::Integer, b::Integer) = ws2812_rgb!(led.sm, r, g, b)
+
+function Base.close(led::WS2812)
+    disable!(led.sm)
+    unclaim!(led.sm)
+end