diff --git a/src/elements/module.rs b/src/elements/module.rs index beb2e7a..3256d1b 100644 --- a/src/elements/module.rs +++ b/src/elements/module.rs @@ -244,18 +244,18 @@ impl Module { } /// Changes the module's start section. - pub fn set_start_section(&mut self, new_start : u32) { - let mut insert_after = 0; - for (i, section) in self.sections_mut().iter_mut().enumerate() { + pub fn set_start_section(&mut self, new_start: u32) { + for section in self.sections_mut().iter_mut() { if let &mut Section::Start(_sect) = section { *section = Section::Start(new_start); return } - if section.id() < 0x8 { - insert_after = i; - } } - self.sections_mut().insert(insert_after + 1, Section::Start(new_start)); + let insert_before = self.sections().iter().enumerate() + .filter_map(|(i, s)| if s.id() > 0x8 { Some(i) } else { None }) + .next() + .unwrap_or(0); + self.sections_mut().insert(insert_before, Section::Start(new_start)); } /// Removes the module's start section. @@ -770,4 +770,19 @@ mod integration_tests { let sections = module.sections().iter().map(|s| s.id()).collect::>(); assert_eq!(sections, vec![1, 2, 3, 6, 7, 8, 9, 10, 11]); } + + #[test] + fn add_start_custom() { + let mut module = deserialize_file("./res/cases/v1/start_add_custom.wasm").expect("failed to deserialize"); + + let sections = module.sections().iter().map(|s| s.id()).collect::>(); + assert_eq!(sections, vec![1, 2, 3, 6, 7, 9, 10, 11, 0]); + + assert!(module.start_section().is_none()); + module.set_start_section(0); + assert_eq!(module.start_section().expect("Did not find any start section"), 0); + + let sections = module.sections().iter().map(|s| s.id()).collect::>(); + assert_eq!(sections, vec![1, 2, 3, 6, 7, 8, 9, 10, 11, 0]); + } }