Skip to content

Commit

Permalink
wip: refactor kaitai-struct-types
Browse files Browse the repository at this point in the history
  • Loading branch information
@Xiphoseer
Xiphoseer committed Sep 16, 2023
1 parent f43c97d commit ce88025
Show file tree
Hide file tree
Showing 16 changed files with 771 additions and 694 deletions.
4 changes: 1 addition & 3 deletions kaitai-codegen-quote/src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -81,8 +81,7 @@ impl Context<'_> {
let doc = spec.doc.as_deref();
let doc_ref = spec.doc_ref.as_ref();
let struct_ty = nc.resolve(key).unwrap();
let st =
structs::codegen_struct(&nc, doc, doc_ref, &spec.seq, &spec.instances, struct_ty)?;
let st = structs::codegen_struct(&nc, doc, doc_ref, &spec.seq, struct_ty).unwrap();
structs.push(st);
}

Expand All @@ -100,7 +99,6 @@ impl Context<'_> {
schema.doc.as_deref(),
schema.doc_ref.as_ref(),
&schema.seq,
&schema.instances,
root_ty,
)
.unwrap();
Expand Down
133 changes: 66 additions & 67 deletions kaitai-codegen-quote/src/parser/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -157,76 +157,67 @@ fn codegen_type_ref_parse(
p_endian: &Ident,
) -> TokenStream {
let in_parent = false;
if let ResolvedTypeKind::Enum(e, w) = &resolved_ty.kind {
let e = nc.get_enum(e).unwrap();
let i_enum = &e.ident;
let q_parser = primitive::wk_parser(w, size, p_endian);
return quote!(::nom::combinator::map_res(#q_parser, #i_enum::try_from));
}
match ty {
TypeRef::WellKnown(w) => primitive::wk_parser(w, size, p_endian),
TypeRef::Named(n) => {
let _named_ty = nc.resolve(n).unwrap();
user_type(nc, _named_ty, self_ty.is_root, in_parent, p_endian)
}
TypeRef::Dynamic { switch_on, cases } => {
match &resolved_ty.kind {
ResolvedTypeKind::Dynamic(_, _) => {
let id = type_ref_id.unwrap();
let fg = self_ty.field_generics.get(id).expect(id);
if fg.external {
fg.parser.to_token_stream()
} else {
variant_parser_expr(nc, self_ty.is_root, fg, in_parent, p_endian)
}
TypeRef::Dynamic { switch_on, cases } => match &resolved_ty.kind {
ResolvedTypeKind::Dynamic(_, _) => {
let id = type_ref_id.unwrap();
let fg = self_ty.field_generics.get(id).expect(id);
if fg.external {
fg.parser.to_token_stream()
} else {
variant_parser_expr(nc, self_ty.is_root, fg, in_parent, p_endian)
}
ResolvedTypeKind::Magic => panic!("Not a TypeRef::Dynamic"),
ResolvedTypeKind::User(_) => panic!("Not a TypeRef::Dynamic"),
ResolvedTypeKind::Enum(_, _) => panic!("Not a TypeRef::Dynamic"),
ResolvedTypeKind::SInt { .. } => panic!("Not a TypeRef::Dynamic"),
ResolvedTypeKind::Float { .. } => panic!("Not a TypeRef::Dynamic"),
ResolvedTypeKind::Str { .. } => panic!("Not a TypeRef::Dynamic"),
ResolvedTypeKind::Bytes { .. } => panic!("Not a TypeRef::Dynamic"),
ResolvedTypeKind::UInt { width, .. } => {
let switch_expr = match switch_on {
}
ResolvedTypeKind::Magic => panic!("Not a TypeRef::Dynamic"),
ResolvedTypeKind::User(_) => panic!("Not a TypeRef::Dynamic"),
ResolvedTypeKind::Enum(_, _) => panic!("Not a TypeRef::Dynamic"),
ResolvedTypeKind::SInt { .. } => panic!("Not a TypeRef::Dynamic"),
ResolvedTypeKind::Float { .. } => panic!("Not a TypeRef::Dynamic"),
ResolvedTypeKind::Str { .. } => panic!("Not a TypeRef::Dynamic"),
ResolvedTypeKind::Bytes { .. } => panic!("Not a TypeRef::Dynamic"),
ResolvedTypeKind::UInt { width, .. } => {
let switch_expr = match switch_on {
AnyScalar::Null => todo!(),
AnyScalar::Bool(_) => todo!(),
AnyScalar::String(expr) => expr::codegen_expr_str(expr),
AnyScalar::UInt(_) => todo!(),
};
let ty = uint_ty(*width);
let mut q_cases = Vec::<TokenStream>::new();
for (k, v) in cases {
let case = match k {
AnyScalar::Bool(b) => quote!(#b),
AnyScalar::Null => todo!(),
AnyScalar::Bool(_) => todo!(),
AnyScalar::String(expr) => expr::codegen_expr_str(expr),
AnyScalar::String(_) => todo!(),
AnyScalar::UInt(_) => todo!(),
};
let ty = uint_ty(*width);
let mut q_cases = Vec::<TokenStream>::new();
for (k, v) in cases {
let case = match k {
AnyScalar::Bool(b) => quote!(#b),
AnyScalar::Null => todo!(),
AnyScalar::String(_) => todo!(),
AnyScalar::UInt(_) => todo!(),
};
let parser = match v {
TypeRef::WellKnown(w) => primitive::wk_parser(w, None, p_endian),
TypeRef::Named(_) => todo!(),
TypeRef::Dynamic { .. } => todo!(),
};
q_cases.push(
quote!(#case => Box::new(::nom::combinator::map(#parser, #ty::from))),
);
}
let exhaustive = true;
if !exhaustive {
q_cases.push(quote!( _ => Box::new(::nom::combinator::map(|_| todo!(), |x: u8| x as #ty))));
}

quote!({
let __parser: Box<dyn FnMut(&'a [u8]) -> ::nom::IResult<&'a[u8], #ty>> = match #switch_expr {
#(#q_cases,)*
};
__parser
})
/**/
let parser = match v {
TypeRef::WellKnown(w) => primitive::wk_parser(w, None, p_endian),
TypeRef::Named(_) => todo!(),
TypeRef::Dynamic { .. } => todo!(),
};
q_cases.push(
quote!(#case => Box::new(::nom::combinator::map(#parser, #ty::from))),
);
}
let exhaustive = true;
if !exhaustive {
q_cases.push(quote!( _ => Box::new(::nom::combinator::map(|_| todo!(), |x: u8| x as #ty))));
}

quote!({
let __parser: Box<dyn FnMut(&'a [u8]) -> ::nom::IResult<&'a[u8], #ty>> = match #switch_expr {
#(#q_cases,)*
};
__parser
})
}
}
},
}
}

Expand All @@ -240,15 +231,23 @@ fn codegen_attr_parse(
) -> TokenStream {
let f_ident = field.ident();
let mut parser = if let Some(ty) = &attr.ty {
codegen_type_ref_parse(
attr.id.as_deref(),
nc,
ty,
attr.size.as_deref(),
self_ty,
field.resolved_ty(),
p_endian,
)
let size = attr.size.as_deref();
if let ResolvedTypeKind::Enum(e, w) = &field.resolved_ty().kind {
let e = nc.get_enum(e).unwrap();
let i_enum = &e.ident;
let q_parser = primitive::wk_parser(w, size, p_endian);
quote!(::nom::combinator::map_res(#q_parser, #i_enum::try_from))
} else {
codegen_type_ref_parse(
attr.id.as_deref(),
nc,
ty,
size,
self_ty,
field.resolved_ty(),
p_endian,
)
}
} else if let Some(contents) = &attr.contents {
let tag = match contents {
Contents::String(s) => quote!(#s),
Expand Down
22 changes: 12 additions & 10 deletions kaitai-codegen-quote/src/parser/primitive.rs
View file
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
use kaitai_struct_types::{EndianSpec, IntTypeRef, WellKnownTypeRef};
use kaitai_struct_types::{EndianSpec, FloatTypeRef, IntTypeRef, WellKnownTypeRef};
use proc_macro2::{Ident, TokenStream};
use quote::quote;

Expand Down Expand Up @@ -85,15 +85,17 @@ pub fn wk_parser(w: &WellKnownTypeRef, size: Option<&str>, p_endian: &Ident) ->
EndianSpec::Big => quote!(::nom::number::complete::be_i64),
},
},
WellKnownTypeRef::F4(e) => match e {
EndianSpec::Implicit => quote!(::nom::number::complete::f32(#p_endian)),
EndianSpec::Little => quote!(::nom::number::complete::le_f32),
EndianSpec::Big => quote!(::nom::number::complete::be_f32),
},
WellKnownTypeRef::F8(e) => match e {
EndianSpec::Implicit => quote!(::nom::number::complete::f64(#p_endian)),
EndianSpec::Little => quote!(::nom::number::complete::le_f64),
EndianSpec::Big => quote!(::nom::number::complete::be_f64),
WellKnownTypeRef::Float(f) => match f {
FloatTypeRef::Float4(e) => match e {
EndianSpec::Implicit => quote!(::nom::number::complete::f32(#p_endian)),
EndianSpec::Little => quote!(::nom::number::complete::le_f32),
EndianSpec::Big => quote!(::nom::number::complete::be_f32),
},
FloatTypeRef::Float8(e) => match e {
EndianSpec::Implicit => quote!(::nom::number::complete::f64(#p_endian)),
EndianSpec::Little => quote!(::nom::number::complete::le_f64),
EndianSpec::Big => quote!(::nom::number::complete::be_f64),
},
},
WellKnownTypeRef::Str => {
if let Some(size_expr) = size {
Expand Down
71 changes: 36 additions & 35 deletions kaitai-codegen-quote/src/structs.rs
View file
Original file line number Diff line number Diff line change
@@ -1,10 +1,7 @@
use kaitai_struct_types::{Attribute, StringOrArray};
use proc_macro2::{Ident, TokenStream};
use quote::{format_ident, quote, ToTokens};
use std::{
collections::{BTreeMap, BTreeSet},
io,
};
use std::collections::BTreeSet;

use crate::{
ctx::NamingContext,
Expand All @@ -14,6 +11,11 @@ use crate::{
},
};

#[derive(Debug)]
pub enum Error {
RequiredFieldGenerics,
}

fn codegen_named_ty(nc: &NamingContext, n: &str) -> TokenStream {
let ty = nc.resolve(n).unwrap();
let mut q_ty = ty.token_stream();
Expand Down Expand Up @@ -66,8 +68,8 @@ fn codegen_cases(
// Field name or enum variant
discriminant: &Ident,
field_generics: Option<&FieldGenerics>,
) -> TokenStream {
let gen = field_generics.unwrap();
) -> Result<TokenStream, Error> {
let gen = field_generics.ok_or(Error::RequiredFieldGenerics)?;
let g = &gen.type_;
let t = &gen.trait_;
if gen.external {
Expand All @@ -84,7 +86,7 @@ fn codegen_cases(
let enclosing_type = gen.var_enum.to_string();
for (i, case_type) in cases.iter().enumerate() {
let n = gen.cases[i].ident();
let inner = codegen_resolved_ty(nc, &case_type.ty, &enclosing_type, tc, n, field_generics);
let inner = codegen_resolved_ty(nc, &case_type.ty, &enclosing_type, tc, n, field_generics)?;
enum_cases.push(quote! {
#n(#inner),
});
Expand Down Expand Up @@ -128,26 +130,25 @@ fn codegen_cases(

#(#trait_impl_cases)*
});
if gen.external {
Ok(if gen.external {
quote!(#g)
} else {
quote!(#v_use)
}
})
}

fn codegen_attr_ty(
nc: &NamingContext,
self_ty: &Type,
field: &Field,
tc: &mut TyContext,
) -> io::Result<TokenStream> {
) -> Result<TokenStream, Error> {
let attr_id = field.ident();
let orig_attr_id = field.id();
let fg = self_ty.field_generics.get(orig_attr_id);
let resolved = field.resolved_ty();
let enclosing_type = self_ty.rust_struct_name.as_str();
let ty = codegen_resolved_ty(nc, resolved, enclosing_type, tc, attr_id, fg);
Ok(ty)
codegen_resolved_ty(nc, resolved, enclosing_type, tc, attr_id, fg)
}

fn codegen_resolved_ty(
Expand All @@ -157,14 +158,14 @@ fn codegen_resolved_ty(
tc: &mut TyContext,
attr_id: &Ident,
fg: Option<&FieldGenerics>,
) -> TokenStream {
let inner = codegen_resolved_ty_kind(&ty.kind, enclosing_type, nc, tc, attr_id, fg);
match &ty.count {
) -> Result<TokenStream, Error> {
let inner = codegen_resolved_ty_kind(&ty.kind, enclosing_type, nc, tc, attr_id, fg)?;
Ok(match &ty.count {
ResolvedTypeCount::Required => inner,
ResolvedTypeCount::Optional => quote!(Option<#inner>),
ResolvedTypeCount::Variable => quote!(Vec<#inner>),
ResolvedTypeCount::Fixed(i) => quote!([#inner; #i]),
}
})
}

fn codegen_resolved_ty_kind(
Expand All @@ -174,18 +175,18 @@ fn codegen_resolved_ty_kind(
tc: &mut TyContext,
attr_id: &Ident,
fg: Option<&FieldGenerics>,
) -> TokenStream {
) -> Result<TokenStream, Error> {
match ty_kind {
ResolvedTypeKind::Dynamic(_switch_on, cases) => {
codegen_cases(cases, nc, tc, enclosing_type, attr_id, fg)
}
ResolvedTypeKind::Magic => quote!(()),
ResolvedTypeKind::User(n) => codegen_named_ty(nc, n),
ResolvedTypeKind::Enum(e, _) => nc.get_enum(e).unwrap().ident.to_token_stream(),
ResolvedTypeKind::UInt { width, .. } => r#type::uint_ty(*width),
ResolvedTypeKind::SInt { width, .. } => r#type::sint_ty(*width),
ResolvedTypeKind::Float { width, .. } => r#type::float_ty(*width),
ResolvedTypeKind::Str { .. } | ResolvedTypeKind::Bytes { .. } => quote!(&'a [u8]),
ResolvedTypeKind::Magic => Ok(quote!(())),
ResolvedTypeKind::User(n) => Ok(codegen_named_ty(nc, n)),
ResolvedTypeKind::Enum(e, _) => Ok(nc.get_enum(e).unwrap().ident.to_token_stream()),
ResolvedTypeKind::UInt { width, .. } => Ok(r#type::uint_ty(*width)),
ResolvedTypeKind::SInt { width, .. } => Ok(r#type::sint_ty(*width)),
ResolvedTypeKind::Float { width, .. } => Ok(r#type::float_ty(*width)),
ResolvedTypeKind::Str { .. } | ResolvedTypeKind::Bytes { .. } => Ok(quote!(&'a [u8])),
}
}

Expand All @@ -195,7 +196,7 @@ fn codegen_attr(
self_ty: &Type,
field: &Field,
tc: &mut TyContext,
) -> io::Result<TokenStream> {
) -> Result<TokenStream, Error> {
let attr_doc = doc::doc_attr(attr);
let serialize_with = parser::serialize_with(attr);

Expand All @@ -213,7 +214,7 @@ fn codegen_struct_body(
self_ty: &Type,
nc: &NamingContext,
tc: &mut TyContext,
) -> Result<TokenStream, io::Error> {
) -> Result<TokenStream, Error> {
if self_ty.is_var_len_str() {
return Ok(quote!((pub &'a [u8]);));
} else if self_ty.is_newtype() {
Expand All @@ -240,14 +241,16 @@ fn codegen_instance_fn(
_nc: &NamingContext,
_self_ty: &Type,
instance: &Field,
_attr: &Attribute,
_tc: &mut TyContext,
) -> io::Result<TokenStream> {
) -> Result<TokenStream, Error> {
let id = format_ident!("parse_{}", instance.ident());
//let ty = codegen_attr_ty(nc, attr, self_ty, instance, tc)?;
let (ty, parser) = match codegen_attr_ty(_nc, _self_ty, instance, _tc) {
Ok(ty) => (ty, quote!(todo!())),
Err(_) => (quote!(()), quote!(todo!())),
};
Ok(quote!(
pub fn #id<'a>(&self, input: &'a [u8]) -> ::nom::IResult<&'a [u8], ()> {
Ok((input, ()))
pub fn #id<'a>(&self, input: &'a [u8]) -> ::nom::IResult<&'a [u8], #ty> {
#parser
}
))
}
Expand All @@ -257,9 +260,8 @@ pub(super) fn codegen_struct(
doc: Option<&str>,
doc_ref: Option<&StringOrArray>,
seq: &[Attribute],
instances: &BTreeMap<String, Attribute>,
self_ty: &Type,
) -> io::Result<TokenStream> {
) -> Result<TokenStream, Error> {
let q_doc = doc::doc_struct(self_ty, nc, doc, doc_ref);
let needs_lifetime = self_ty.needs_lifetime;

Expand All @@ -275,8 +277,7 @@ pub(super) fn codegen_struct(
let q_parser = parser::codegen_parser_fn(self_ty, seq, &tc, nc);

for instance in &self_ty.instances {
let attr = instances.get(instance.id()).unwrap();
let _fn = codegen_instance_fn(nc, self_ty, instance, attr, &mut tc)?;
let _fn = codegen_instance_fn(nc, self_ty, instance, &mut tc)?;
q_impls.push(_fn)
}

Expand Down
Loading

0 comments on commit ce88025

Please sign in to comment.