bitstruct.rs

#![feature(macro_registrar, managed_boxes, quote)]
#![crate_type = "dylib"]

extern crate syntax;

use syntax::codemap;
use syntax::ast;
use syntax::ext::base;
use syntax::ext::build::AstBuilder;
use syntax::parse::token;
use syntax::util::small_vector::SmallVector;

// Entry point
#[macro_registrar]
pub fn macro_registrar(register: |ast::Name, base::SyntaxExtension|) {
  register(
    token::intern("bitstruct"),
    base::IdentTT(
      ~base::BasicIdentMacroExpander {
        expander: expand_syntax_ext,
        span: None,
      },
      None));
}

// Struct for returning multiple Item's as a ~MacResult
pub struct MyMacItems {
  i: Vec<@ast::Item>
}
impl MyMacItems {
  pub fn new(i: Vec<@ast::Item>) -> ~base::MacResult {
    ~MyMacItems { i: i } as ~base::MacResult
  }
}
impl base::MacResult for MyMacItems {
  fn make_items(&self) -> Option<SmallVector<@ast::Item>> {
    Some(SmallVector::many(self.i.clone()))
  }
}

fn expand_syntax_ext(cx: &mut base::ExtCtxt, sp: codemap::Span, struct_name: ast::Ident, tts: Vec<ast::TokenTree>) -> ~base::MacResult {
  // Contains to resulting Items to be returned as a MacResult
  let mut res = Vec::new();

  // Get the fields of the struct
  let fields = get_fields(cx, sp, tts);

  // The total number of bits
  let mut total_bits: uint = 0;

  // Count the total number of bits
  for &(ref _name, _prot, _sign, bits) in fields.iter() {
    total_bits += bits;
  }

  total_bits = ((total_bits - 1) | 7) + 1;

  let total_bytes: uint = total_bits / 8;

  // Generate the actual struct
  res.push(
    quote_item!(
      &*cx,
      pub struct $struct_name {
        pub bytes: [u8, ..$total_bytes]
      }
    ).unwrap()
  );

  // Generate a Default instance
  res.push(
    quote_item!(
      &*cx,
      impl std::default::Default for $struct_name {
        fn default() -> $struct_name {
          $struct_name { bytes: [0, ..$total_bytes] }
        }
      }
    ).unwrap()
  );


  // Total number of bits added so far
  let mut cur_offset: uint = 0;

  // Generate the impls for getting and setting fields
  for &(ref name, prot, sign, bits) in fields.iter() {
    if name.get().char_at(0) != '_' {
      // Generate the actual getter and setter
      res.push(
        generate_methods(
          cx,
          sp,
          struct_name,
          name.get(),
          prot,
          sign,
          cur_offset,
          bits,
        )
      );
    }

    // Update the offset
    cur_offset += bits;
  }

  MyMacItems::new(res)
}

fn generate_methods(
  cx: &mut base::ExtCtxt,
  sp: codemap::Span,
  struct_name: ast::Ident,
  name: &str,
  prot: ast::Visibility,
  sign: bool,
  first_bit: uint,
  size_bits: uint,
) -> @ast::Item {

  // ast::Visibility does not have a to_tokens, so we work around that
  let prot = match prot {
    ast::Public => quote_tokens!(&*cx, pub),
    ast::Inherited => Vec::new(),
  };

  // Create the type and get a variable for the type size
  let (ty_bits, ty): (uint, @ast::Ty) = match (sign, size_bits) {
    (true,  1..8)   => (8,  quote_ty!(&*cx, i8)),
    (true,  9..16)  => (16, quote_ty!(&*cx, i16)),
    (true,  17..32) => (32, quote_ty!(&*cx, i32)),
    (true,  33..64) => (64, quote_ty!(&*cx, i64)),
    (false, 1..8)   => (8,  quote_ty!(&*cx, u8)),
    (false, 9..16)  => (16, quote_ty!(&*cx, u16)),
    (false, 17..32) => (32, quote_ty!(&*cx, u32)),
    (false, 33..64) => (64, quote_ty!(&*cx, u64)),
    _ => cx.span_fatal(sp, format!("bitstruct does support bitsize {}", size_bits)),
  };

  let skew               = first_bit % 8;
  let first_byte         = first_bit / 8;
  let last_bit           = first_bit + size_bits - 1;
  let last_byte          = last_bit / 8;
  let used_bytes         = 1 + last_byte - first_byte;
  let get                = token::str_to_ident("get_" + name);
  let set                = token::str_to_ident("set_" + name);

  let mask0: u64 = if size_bits == 64 { 0 } else { -1 << size_bits };
  let mask0: u64 = mask0 | ((1 << skew) - 1);
  let mask1: u64 = !((1 << skew) - 1);

  (quote_item!(
    &*cx,
    impl $struct_name {
      $prot fn $get(&self) -> $ty {
        let mut arr: [u64, ..2] = [0, 0];
        unsafe {
          std::ptr::copy_nonoverlapping_memory::<u8>(
            std::cast::transmute(arr.as_mut_ptr()),
            self.bytes.as_ptr().offset($first_byte as int),
            $used_bytes
          );
        }
        let val: u64 = if $skew == 0 { arr[0] } else { (arr[0] >> $skew) | (arr[1] << (64 - $skew)) };
        let val: $ty = val as $ty;
        let val = val << ($ty_bits - $size_bits);
        let val = val >> ($ty_bits - $size_bits);
        val
      }

      $prot fn $set(&mut self, val: $ty) {
        let mut arr: [u64, ..2] = [0, 0];
        unsafe {
          std::ptr::copy_nonoverlapping_memory::<u8>(
            std::cast::transmute(arr.as_mut_ptr()),
            self.bytes.as_ptr().offset($first_byte as int),
            $used_bytes
          );
        }

        let val = val as u64;
        arr[0] = (arr[0] & $mask0) | ((val << $skew) & !$mask0);
        arr[1] = (arr[1] & $mask1) | ((val >> (64 - $skew)) & !$mask1);

        unsafe {
          std::ptr::copy_nonoverlapping_memory::<u8>(
            self.bytes.as_mut_ptr().offset($first_byte as int),
            std::cast::transmute(arr.as_ptr()),
            $used_bytes
          );
        }
      }
    }
  )).unwrap()
}

fn get_fields(cx: &mut base::ExtCtxt, sp: codemap::Span, tts: Vec<ast::TokenTree>) -> Vec<(token::InternedString, ast::Visibility, bool, uint)> {
  // Parse the tts to a struct
  let parsed = parse_tts(cx, sp, tts);

  // Return value to be populated
  let mut fields = Vec::new();

  // Iterate over the fields to get the relevant data (while doing sanity checking)
  for field in parsed.fields.iter() {
    let field = &field.node;
    let (name, prot) = match field.kind {
      // Get name and protection. The catch-all case should never be hit
      ast::NamedField(n, p) => (token::get_ident(n), p),
      _ => cx.span_fatal(sp, "bitstruct only support named fields")
    };

    // No attributes please
    if !field.attrs.is_empty() {
      cx.span_fatal(sp, "bitstruct does not support attributes");
    }

    // Get the path of the type
    let path = match field.ty.node {
      ast::TyPath(ref ps, _, _) => ps,
      _ => cx.span_fatal(sp, "bitstruct only supports types of format [ui][0-9]+")
    };

    // Only non-pathed types
    if path.global || path.segments.len() != 1 {
      cx.span_fatal(sp, "bitstruct does not support pathed types+")
    }

    // Get the first (and only) segment
    let segment = path.segments.get(0);

    // No parameters please
    if !segment.lifetimes.is_empty() || segment.types.len() != 0 {
      cx.span_fatal(sp, "bitstruct does not support type parameters")
    }

    // Get the actual type
    let ident = token::get_ident(segment.identifier);
    let typ: &str = ident.get();

    // "Parse" the type
    let (sign, bits) = match (typ[0], from_str(typ.slice_from(1))) {
      (105, Some(n)) => (true,  n), // 'i'
      (117, Some(n)) => (false, n), // 'u'
      _ => cx.span_fatal(sp, "bitstruct only supports types of format [ui][0-9]+")
    };

    // Add the data to the result
    fields.push((name, prot, sign, bits));
  }

  fields
}

fn parse_tts(cx: &mut base::ExtCtxt, sp: codemap::Span, tts: Vec<ast::TokenTree>) -> @ast::StructDef {
  // Parse the tts after adding the struct-cruff
  let parsed = quote_item!(&*cx, struct some_struct { $tts });

  // Was it parsed correctly?
  let item = parsed.unwrap_or_else(
    || cx.span_fatal(sp, "Could not parse struct-like definition")
  );

  // Get out the actual StructDef
  match item.node {
    ast::ItemStruct(s, _) => s,
    _ => unreachable!()
  }
}