Get the application heap size from the environment

Currently, the application heap size is hard coded in entry_point/mod.rs.
If the app heap size needs to be adjusted, we have to manually change
the code and then updat elf2tab as well. Instead of relying on this
potentially troublesome process, allow taking the value from the environment
via APP_HEAP_SIZE. This is parsed at compile time.
diff --git a/build.rs b/build.rs
index 7160086..502c779 100644
--- a/build.rs
+++ b/build.rs
@@ -9,12 +9,13 @@
 fn main() {
     static ENV_VAR: &str = "PLATFORM";
     static FILE_NAME: &str = "platform";
+    static APP_HEAP_SIZE: &str = "APP_HEAP_SIZE";
 
     println!("cargo:rerun-if-env-changed={}", ENV_VAR);
+    println!("cargo:rerun-if-env-changed={}", APP_HEAP_SIZE);
     println!("cargo:rerun-if-changed={}", FILE_NAME);
 
-    let platform_name =
-        read_board_name_from_env_var(ENV_VAR).or_else(|| read_board_name_from_file(FILE_NAME));
+    let platform_name = read_env_var(ENV_VAR).or_else(|| read_board_name_from_file(FILE_NAME));
     if let Some(platform_name) = platform_name {
         println!("cargo:rustc-env={}={}", ENV_VAR, platform_name);
         copy_linker_file(&platform_name.trim());
@@ -24,9 +25,16 @@
              Remember to manually specify a linker file.",
         );
     }
+
+    if let Some(s) = read_env_var(APP_HEAP_SIZE) {
+        println!("cargo:rustc-env=APP_HEAP_SIZE={}", s);
+    } else {
+        // Just use a default of 1024 if nothing is passed in
+        println!("cargo:rustc-env=APP_HEAP_SIZE=1024");
+    }
 }
 
-fn read_board_name_from_env_var(env_var: &str) -> Option<String> {
+fn read_env_var(env_var: &str) -> Option<String> {
     env::var_os(env_var).map(|os_string| os_string.into_string().unwrap())
 }
 
diff --git a/codegen/src/lib.rs b/codegen/src/lib.rs
index 9804e29..ccb0bfb 100644
--- a/codegen/src/lib.rs
+++ b/codegen/src/lib.rs
@@ -9,6 +9,55 @@
 use syn::ItemFn;
 
 #[allow(clippy::needless_doctest_main)]
+/// Procedural macro. This generates a function to read and parse
+/// an environment variable to a usize at compile time. We have to use a
+/// function since procedural macros can't generate expressions at the moment.
+/// Example:
+///
+///
+/// ```ignore
+/// make_read_env_var!("MAX_COUNT")
+///
+/// fn do_it() {
+///     let max = read_MAX_COUNT();
+/// }
+/// ```
+#[proc_macro]
+pub fn make_read_env_var(input: TokenStream) -> TokenStream {
+    let env_arg = syn::parse_macro_input!(input as syn::LitStr);
+    let env_var_name = env_arg.value();
+    let app_size = match std::env::var(&env_var_name) {
+        Ok(r) => r,
+        Err(_r) => {
+            return syn::Error::new(
+                env_arg.span(),
+                format!("Failed to find {} in environment, is it set?", env_var_name),
+            )
+            .to_compile_error()
+            .into()
+        }
+    };
+    let size = match app_size.parse::<usize>() {
+        Ok(r) => r,
+        Err(_r) => {
+            return syn::Error::new(
+                env_arg.span(),
+                format!("Environment var {} can't be parsed to usize", env_var_name),
+            )
+            .to_compile_error()
+            .into()
+        }
+    };
+
+    let concat = format!("read_{}", env_var_name);
+    let i = syn::Ident::new(&concat, proc_macro2::Span::call_site());
+    let result = quote! {
+        fn #i() -> usize { #size }
+    };
+    result.into()
+}
+
+#[allow(clippy::needless_doctest_main)]
 /// Procedural attribute macro. This is meant to be applied to a binary's async
 /// `main()`, transforming into a function that returns a type acceptable for
 /// `main()`. In other words, this will not compile with libtock-rs:
diff --git a/core/Cargo.toml b/core/Cargo.toml
index 64fa96e..3110346 100644
--- a/core/Cargo.toml
+++ b/core/Cargo.toml
@@ -11,3 +11,4 @@
 
 [dependencies]
 linked_list_allocator = { optional = true, version = "=0.6.5", default-features = false }
+libtock_codegen = { path = "../codegen" }
diff --git a/core/build.rs b/core/build.rs
new file mode 100644
index 0000000..e91df46
--- /dev/null
+++ b/core/build.rs
@@ -0,0 +1,22 @@
+use std::env;
+
+fn main() {
+    static APP_HEAP_SIZE: &str = "APP_HEAP_SIZE";
+
+    println!("cargo:rerun-if-env-changed={}", APP_HEAP_SIZE);
+
+    set_default_env(APP_HEAP_SIZE, "1024");
+}
+
+fn set_default_env(env_var: &str, default: &str) {
+    if let Some(s) = read_env_var(env_var) {
+        println!("cargo:rustc-env={}={}", env_var, s);
+    } else {
+        // Just use a default of 1024 if nothing is passed in
+        println!("cargo:rustc-env={}={}", env_var, default);
+    }
+}
+
+fn read_env_var(env_var: &str) -> Option<String> {
+    env::var_os(env_var).map(|os_string| os_string.into_string().unwrap())
+}
diff --git a/core/src/entry_point/mod.rs b/core/src/entry_point/mod.rs
index d63b174..c2628e1 100644
--- a/core/src/entry_point/mod.rs
+++ b/core/src/entry_point/mod.rs
@@ -75,6 +75,9 @@
     stack_size: usize,
 }
 
+//Procedural macro to generate a function to read APP_HEAP_SIZE
+libtock_codegen::make_read_env_var!("APP_HEAP_SIZE");
+
 /// Rust setup, called by _start. Uses the extern "C" calling convention so that
 /// the assembly in _start knows how to call it (the Rust ABI is not defined).
 /// Sets up the data segment (including relocations) and the heap, then calls
@@ -116,21 +119,19 @@
     // `sbrk` system call to dynamically request heap memory from the kernel, we
     // need to tell `linked_list_allocator` where the heap starts and ends.
     //
-    // Heap size is set using `elf2tab` with `--app-heap` option, which is
-    // currently at 1024. If you change the `elf2tab` heap size, make sure to
-    // make the corresponding change here.
-    const HEAP_SIZE: usize = 1024;
+    // We get this from the environment to make it easy to set per compile.
+    let app_heap_size: usize = read_APP_HEAP_SIZE();
 
     // Make the heap start exactly at bss_end. The suggested _app_heap_break
     // is almost always going to be too big and leads to us wasting memory.
     let app_heap_start = bss_end;
-    let app_heap_end = app_heap_start + HEAP_SIZE;
+    let app_heap_end = app_heap_start + app_heap_size;
 
     // Tell the kernel the new app heap break.
     memop::set_brk(app_heap_end as *const u8);
 
     #[cfg(feature = "alloc")]
-    crate::alloc::HEAP.init(app_heap_start, HEAP_SIZE);
+    crate::alloc::HEAP.init(app_heap_start, app_heap_size);
 
     main(0, ptr::null());
 
diff --git a/tools/flash.sh b/tools/flash.sh
index 7d20190..7d5dbbc 100755
--- a/tools/flash.sh
+++ b/tools/flash.sh
@@ -4,6 +4,10 @@
 
 artifact="$(basename $1)"
 rust_target_folder="$(readlink -f $(dirname $1)/../..)"
+if [ -z $APP_HEAP_SIZE ]; then
+	echo "Set APP_HEAP_SIZE to a value"
+	exit 1
+fi
 
 case "${PLATFORM}" in
     "nrf52"|"nrf52840")
@@ -39,7 +43,7 @@
 mkdir -p "${libtock_target_path}"
 cp "$1" "${elf_file_name}"
 
-elf2tab -n "${artifact}" -o "${tab_file_name}" "${elf_file_name}" --stack 2048 --app-heap 1024 --kernel-heap 1024 --protected-region-size=64
+elf2tab -n "${artifact}" -o "${tab_file_name}" "${elf_file_name}" --stack 2048 --app-heap $APP_HEAP_SIZE --kernel-heap 1024 --protected-region-size=64
 
 if [ $tockload == "n" ]; then
 	echo "Skipping flashing for platform \"${PLATFORM}\""