1#!/usr/bin/env python3
2
3"""
4Purpose
5
6This script is for comparing the size of the library files from two
7different Git revisions within an Mbed TLS repository.
8The results of the comparison is formatted as csv and stored at a
9configurable location.
10Note: must be run from Mbed TLS root.
11"""
12
13# Copyright The Mbed TLS Contributors
14# SPDX-License-Identifier: Apache-2.0
15#
16# Licensed under the Apache License, Version 2.0 (the "License"); you may
17# not use this file except in compliance with the License.
18# You may obtain a copy of the License at
19#
20# http://www.apache.org/licenses/LICENSE-2.0
21#
22# Unless required by applicable law or agreed to in writing, software
23# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
24# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
25# See the License for the specific language governing permissions and
26# limitations under the License.
27
28import argparse
29import os
30import subprocess
31import sys
32
33class CodeSizeComparison:
34    """Compare code size between two Git revisions."""
35
36    def __init__(self, old_revision, new_revision, result_dir):
37        """
38        old_revision: revision to compare against
39        new_revision:
40        result_dir: directory for comparision result
41        """
42        self.repo_path = "."
43        self.result_dir = os.path.abspath(result_dir)
44        os.makedirs(self.result_dir, exist_ok=True)
45
46        self.csv_dir = os.path.abspath("code_size_records/")
47        os.makedirs(self.csv_dir, exist_ok=True)
48
49        self.old_rev = old_revision
50        self.new_rev = new_revision
51        self.git_command = "git"
52        self.make_command = "make"
53
54    @staticmethod
55    def check_repo_path():
56        if not all(os.path.isdir(d) for d in ["include", "library", "tests"]):
57            raise Exception("Must be run from Mbed TLS root")
58
59    @staticmethod
60    def validate_revision(revision):
61        result = subprocess.check_output(["git", "rev-parse", "--verify",
62                                          revision + "^{commit}"], shell=False)
63        return result
64
65    def _create_git_worktree(self, revision):
66        """Make a separate worktree for revision.
67        Do not modify the current worktree."""
68
69        if revision == "current":
70            print("Using current work directory.")
71            git_worktree_path = self.repo_path
72        else:
73            print("Creating git worktree for", revision)
74            git_worktree_path = os.path.join(self.repo_path, "temp-" + revision)
75            subprocess.check_output(
76                [self.git_command, "worktree", "add", "--detach",
77                 git_worktree_path, revision], cwd=self.repo_path,
78                stderr=subprocess.STDOUT
79            )
80        return git_worktree_path
81
82    def _build_libraries(self, git_worktree_path):
83        """Build libraries in the specified worktree."""
84
85        my_environment = os.environ.copy()
86        subprocess.check_output(
87            [self.make_command, "-j", "lib"], env=my_environment,
88            cwd=git_worktree_path, stderr=subprocess.STDOUT,
89        )
90
91    def _gen_code_size_csv(self, revision, git_worktree_path):
92        """Generate code size csv file."""
93
94        csv_fname = revision + ".csv"
95        if revision == "current":
96            print("Measuring code size in current work directory.")
97        else:
98            print("Measuring code size for", revision)
99        result = subprocess.check_output(
100            ["size library/*.o"], cwd=git_worktree_path, shell=True
101        )
102        size_text = result.decode()
103        csv_file = open(os.path.join(self.csv_dir, csv_fname), "w")
104        for line in size_text.splitlines()[1:]:
105            data = line.split()
106            csv_file.write("{}, {}\n".format(data[5], data[3]))
107
108    def _remove_worktree(self, git_worktree_path):
109        """Remove temporary worktree."""
110        if git_worktree_path != self.repo_path:
111            print("Removing temporary worktree", git_worktree_path)
112            subprocess.check_output(
113                [self.git_command, "worktree", "remove", "--force",
114                 git_worktree_path], cwd=self.repo_path,
115                stderr=subprocess.STDOUT
116            )
117
118    def _get_code_size_for_rev(self, revision):
119        """Generate code size csv file for the specified git revision."""
120
121        # Check if the corresponding record exists
122        csv_fname = revision + ".csv"
123        if (revision != "current") and \
124           os.path.exists(os.path.join(self.csv_dir, csv_fname)):
125            print("Code size csv file for", revision, "already exists.")
126        else:
127            git_worktree_path = self._create_git_worktree(revision)
128            self._build_libraries(git_worktree_path)
129            self._gen_code_size_csv(revision, git_worktree_path)
130            self._remove_worktree(git_worktree_path)
131
132    def compare_code_size(self):
133        """Generate results of the size changes between two revisions,
134        old and new. Measured code size results of these two revisions
135        must be available."""
136
137        old_file = open(os.path.join(self.csv_dir, self.old_rev + ".csv"), "r")
138        new_file = open(os.path.join(self.csv_dir, self.new_rev + ".csv"), "r")
139        res_file = open(os.path.join(self.result_dir, "compare-" + self.old_rev
140                                     + "-" + self.new_rev + ".csv"), "w")
141
142        res_file.write("file_name, this_size, old_size, change, change %\n")
143        print("Generating comparision results.")
144
145        old_ds = {}
146        for line in old_file.readlines()[1:]:
147            cols = line.split(", ")
148            fname = cols[0]
149            size = int(cols[1])
150            if size != 0:
151                old_ds[fname] = size
152
153        new_ds = {}
154        for line in new_file.readlines()[1:]:
155            cols = line.split(", ")
156            fname = cols[0]
157            size = int(cols[1])
158            new_ds[fname] = size
159
160        for fname in new_ds:
161            this_size = new_ds[fname]
162            if fname in old_ds:
163                old_size = old_ds[fname]
164                change = this_size - old_size
165                change_pct = change / old_size
166                res_file.write("{}, {}, {}, {}, {:.2%}\n".format(fname, \
167                               this_size, old_size, change, float(change_pct)))
168            else:
169                res_file.write("{}, {}\n".format(fname, this_size))
170        return 0
171
172    def get_comparision_results(self):
173        """Compare size of library/*.o between self.old_rev and self.new_rev,
174        and generate the result file."""
175        self.check_repo_path()
176        self._get_code_size_for_rev(self.old_rev)
177        self._get_code_size_for_rev(self.new_rev)
178        return self.compare_code_size()
179
180def main():
181    parser = argparse.ArgumentParser(
182        description=(
183            """This script is for comparing the size of the library files
184            from two different Git revisions within an Mbed TLS repository.
185            The results of the comparison is formatted as csv, and stored at
186            a configurable location.
187            Note: must be run from Mbed TLS root."""
188        )
189    )
190    parser.add_argument(
191        "-r", "--result-dir", type=str, default="comparison",
192        help="directory where comparison result is stored, \
193              default is comparison",
194    )
195    parser.add_argument(
196        "-o", "--old-rev", type=str, help="old revision for comparison.",
197        required=True,
198    )
199    parser.add_argument(
200        "-n", "--new-rev", type=str, default=None,
201        help="new revision for comparison, default is the current work \
202              directory, including uncommited changes."
203    )
204    comp_args = parser.parse_args()
205
206    if os.path.isfile(comp_args.result_dir):
207        print("Error: {} is not a directory".format(comp_args.result_dir))
208        parser.exit()
209
210    validate_res = CodeSizeComparison.validate_revision(comp_args.old_rev)
211    old_revision = validate_res.decode().replace("\n", "")
212
213    if comp_args.new_rev is not None:
214        validate_res = CodeSizeComparison.validate_revision(comp_args.new_rev)
215        new_revision = validate_res.decode().replace("\n", "")
216    else:
217        new_revision = "current"
218
219    result_dir = comp_args.result_dir
220    size_compare = CodeSizeComparison(old_revision, new_revision, result_dir)
221    return_code = size_compare.get_comparision_results()
222    sys.exit(return_code)
223
224
225if __name__ == "__main__":
226    main()
227