From 98570fc1d7cb98d4b86496c0024130804677b173 Mon Sep 17 00:00:00 2001
From: Sergey Pepyakin <sergei@parity.io>
Date: Tue, 16 Apr 2019 16:14:10 +0200
Subject: [PATCH] Estimate capacity.

---
 src/prepare/compile.rs | 7 ++++---
 validation/src/func.rs | 2 +-
 validation/src/lib.rs  | 9 ++++++---
 3 files changed, 11 insertions(+), 7 deletions(-)

diff --git a/src/prepare/compile.rs b/src/prepare/compile.rs
index 43869ec..6714a18 100644
--- a/src/prepare/compile.rs
+++ b/src/prepare/compile.rs
@@ -1,7 +1,7 @@
 #[allow(unused_imports)]
 use alloc::prelude::v1::*;
 
-use parity_wasm::elements::{BlockType, Instruction};
+use parity_wasm::elements::{BlockType, Instruction, FuncBody};
 
 use validation::func::{
     require_label, top_label, BlockFrame, FunctionValidationContext, StackValueType, StartedWith,
@@ -69,9 +69,10 @@ pub struct Compiler {
 
 impl FuncValidator for Compiler {
     type Output = isa::Instructions;
-    fn new(_module: &FunctionValidationContext) -> Self {
+    fn new(_ctx: &FunctionValidationContext, body: &FuncBody) -> Self {
+        let code_len = body.code().elements().len();
         let mut compiler = Compiler {
-            sink: Sink::with_instruction_capacity(0), // TODO: Estimate instruction count.
+            sink: Sink::with_instruction_capacity(code_len),
             label_stack: Vec::new(),
         };
 
diff --git a/validation/src/func.rs b/validation/src/func.rs
index 0976675..2f0f742 100644
--- a/validation/src/func.rs
+++ b/validation/src/func.rs
@@ -120,7 +120,7 @@ pub fn drive<T: FuncValidator>(
         result_ty,
     );
 
-    let mut validator = T::new(&context);
+    let mut validator = T::new(&context, body);
 
     for (position, instruction) in code.iter().enumerate() {
         validator
diff --git a/validation/src/lib.rs b/validation/src/lib.rs
index d8bf94b..11fa542 100644
--- a/validation/src/lib.rs
+++ b/validation/src/lib.rs
@@ -35,7 +35,7 @@ use std::collections::HashSet;
 use self::context::ModuleContextBuilder;
 use parity_wasm::elements::{
     BlockType, External, GlobalEntry, GlobalType, InitExpr, Instruction, Internal, MemoryType,
-    Module, ResizableLimits, TableType, Type, ValueType,
+    Module, ResizableLimits, TableType, Type, ValueType, FuncBody,
 };
 
 pub mod context;
@@ -83,7 +83,10 @@ pub trait Validator {
 
 pub trait FuncValidator {
     type Output;
-    fn new(ctx: &func::FunctionValidationContext) -> Self;
+    fn new(
+        ctx: &func::FunctionValidationContext,
+        body: &FuncBody,
+    ) -> Self;
     fn next_instruction(
         &mut self,
         ctx: &mut func::FunctionValidationContext,
@@ -119,7 +122,7 @@ pub struct PlainFuncValidator;
 impl FuncValidator for PlainFuncValidator {
     type Output = ();
 
-    fn new(_ctx: &func::FunctionValidationContext) -> PlainFuncValidator {
+    fn new(_ctx: &func::FunctionValidationContext, _body: &FuncBody) -> PlainFuncValidator {
         PlainFuncValidator
     }