From 3612442adc2e0ea0f1ea92652997ec0b3b76722c Mon Sep 17 00:00:00 2001
From: "Hui, Chunyang" <sanqian.hcy@antfin.com>
Date: Wed, 21 Jul 2021 07:12:54 +0000
Subject: [PATCH] Fix an error when calculating elf memory usage

VMLayout was mistakenly used to calculate the memory usage. This
commit is to fix this and seperate VMLayout "add" and "extend"
methods.
---
 src/libos/src/vm/process_vm.rs |  2 +-
 src/libos/src/vm/vm_layout.rs  | 17 ++++++++++++++++-
 2 files changed, 17 insertions(+), 2 deletions(-)

diff --git a/src/libos/src/vm/process_vm.rs b/src/libos/src/vm/process_vm.rs
index 0be3b85d..3887ef78 100644
--- a/src/libos/src/vm/process_vm.rs
+++ b/src/libos/src/vm/process_vm.rs
@@ -101,7 +101,7 @@ impl<'a, 'b> ProcessVMBuilder<'a, 'b> {
         let process_layout = elf_layouts.iter().chain(other_layouts.iter()).fold(
             VMLayout::new_empty(),
             |mut process_layout, sub_layout| {
-                process_layout.extend(&sub_layout);
+                process_layout.add(&sub_layout);
                 process_layout
             },
         );
diff --git a/src/libos/src/vm/vm_layout.rs b/src/libos/src/vm/vm_layout.rs
index d025f038..578fa5c1 100644
--- a/src/libos/src/vm/vm_layout.rs
+++ b/src/libos/src/vm/vm_layout.rs
@@ -21,7 +21,8 @@ impl VMLayout {
         }
     }
 
-    pub fn extend(&mut self, more_space: &VMLayout) -> &mut Self {
+    // This is used to add "more_space" to VM layout
+    pub fn add(&mut self, more_space: &VMLayout) -> &mut Self {
         if more_space.size == 0 {
             return self;
         }
@@ -31,6 +32,20 @@ impl VMLayout {
         self
     }
 
+    // This is used to get the bigger and aligned VM layout
+    pub fn extend(&mut self, new_space: &VMLayout) -> &mut Self {
+        if new_space.size == 0 {
+            return self;
+        }
+
+        self.align = max(self.align, new_space.align);
+        self.size = {
+            let size = max(self.size, new_space.size);
+            align_up(size, self.align)
+        };
+        self
+    }
+
     pub fn size(&self) -> usize {
         self.size
     }