coap-message 0.1.0

//! This module provides [HeapMessage], an implementation of all the message traits backed by heap
//! memory.
//!
//! It serves both to easily get an owned copy of a message (using [`HeapMessage::new()`] and
//! [`MinimalWritableMessage::set_from_message()`]), especially if that needs frequent manipulation
//! of options, and as an example of how to the message traits can be implemented.
//!
//! This module is available only if the `alloc` feature is enabled on this crate.

use crate::message::*;

extern crate alloc;
use alloc::{vec, vec::Vec};

/// A heap CoAP message backed by allocated memory
///
/// It stores its payload in a [`Vec`], and uses a [`BTreeMap<_, Vec<_>>`](alloc::collections::BTreeMap) to
/// store all the individual option values.
///
/// It offers a few methods for direct manipulation of options, even out of sequence, that can not
/// be expected from a general message buffer and are thus not captured in traits.
#[derive(Debug)]
pub struct HeapMessage {
    code: u8,
    options: alloc::collections::BTreeMap<u16, Vec<Vec<u8>>>,
    payload: Vec<u8>,
}

impl HeapMessage {
    pub fn new() -> Self {
        Self {
            code: 0,
            options: alloc::collections::BTreeMap::new(),
            payload: vec![],
        }
    }

    // The following are not used via the WritableMessage as they're really custom to a
    // GUI-manager message

    /// Replace the occurrence'th value of the optnum option in the message
    ///
    /// Panics if there's not ``occurrence + 1`` options of that number.
    pub fn change_option(&mut self, optnum: u16, occurrence: usize, value: impl Into<Vec<u8>>) {
        self.options.get_mut(&optnum).unwrap()[occurrence] = value.into();
    }

    /// Remove the occurrence'th option of option number optnum
    ///
    /// Panics if there's not ``occurrence + 1`` options of that number.
    pub fn remove_option(&mut self, optnum: u16, occurrence: usize) {
        let vec = self.options.get_mut(&optnum).unwrap();
        vec.remove(occurrence);
        if vec.len() == 0 {
            self.options.remove(&optnum);
        }
    }

    /// Like MinimalWritableMessage::add_option, but allowing arbitrary access
    pub fn add_option(&mut self, optnum: u16, data: &[u8]) {
        self.options
            .entry(optnum)
            .or_insert(vec![])
            .push(data.to_vec());
    }
}

impl MinimalWritableMessage for HeapMessage {
    type Code = u8;
    type OptionNumber = u16;
    fn set_code(&mut self, code: u8) {
        self.code = code;
    }
    fn add_option(&mut self, optnum: u16, data: &[u8]) {
        // Only a debug-assert here because it's not something this needs to enforce, but more a
        // helper flag code that abuses an impl MinimalWritableMessage even when backed by the
        // HeapMessage that doesn't care
        debug_assert!(
            // Can be simplified once https://github.com/rust-lang/rust/issues/62924 is closed
            &optnum >= self.options.keys().max().unwrap_or(&0),
            "Option appended out-of-order");

        self.add_option(optnum, data);
    }
    fn set_payload(&mut self, payload: &[u8]) {
        self.payload = payload.to_vec();
    }
}

pub struct MessageOption<'a> {
    number: u16,
    value: &'a [u8],
}

impl<'a> crate::message::MessageOption for MessageOption<'a> {
    fn number(&self) -> u16 { self.number }
    fn value(&self) -> &[u8] { self.value }
}

pub struct ReadCursor<'a> {
    iter: alloc::collections::btree_map::Iter<'a, u16, Vec<Vec<u8>>>,
    popped: (u16, &'a [Vec<u8>]),
}

impl<'a> Iterator for ReadCursor<'a> {
    type Item = MessageOption<'a>;
    fn next(&mut self) -> Option<MessageOption<'a>> {
        if self.popped.1.len() == 0 {
            self.popped = self.iter.next()
                .map(|(k, v)| (*k, v.as_slice()))
                ?;

            // Avoided by construction
            debug_assert!(self.popped.1.len() > 0, "HeapMessage had present but empty option");
        }

        let (first, rest) = self.popped.1.split_at(1);
        self.popped.1 = rest;
        Some(MessageOption { number: self.popped.0, value: first[0].as_slice() })
    }
}

impl<'a> WithSortedOptions<'a> for HeapMessage {
    // By virtue of the order of options used in BTreeMap
}

impl<'m> ReadableMessage<'m> for HeapMessage {
    type Code = u8;
    type MessageOption = MessageOption<'m>;
    type OptionsIter = ReadCursor<'m>;

    fn code(&self) -> u8 { self.code }
    fn payload(&self) -> &[u8] { &self.payload }
    fn options(&'m self) -> ReadCursor<'m> {
        ReadCursor { iter: self.options.iter(), popped: (0, &[]) }
    }
}