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
33from mbedtls_dev import build_tree
34
35
36class CodeSizeComparison:
37    """Compare code size between two Git revisions."""
38
39    def __init__(self, old_revision, new_revision, result_dir):
40        """
41        old_revision: revision to compare against
42        new_revision:
43        result_dir: directory for comparison result
44        """
45        self.repo_path = "."
46        self.result_dir = os.path.abspath(result_dir)
47        os.makedirs(self.result_dir, exist_ok=True)
48
49        self.csv_dir = os.path.abspath("code_size_records/")
50        os.makedirs(self.csv_dir, exist_ok=True)
51
52        self.old_rev = old_revision
53        self.new_rev = new_revision
54        self.git_command = "git"
55        self.make_command = "make"
56
57    @staticmethod
58    def validate_revision(revision):
59        result = subprocess.check_output(["git", "rev-parse", "--verify",
60                                          revision + "^{commit}"], shell=False)
61        return result
62
63    def _create_git_worktree(self, revision):
64        """Make a separate worktree for revision.
65        Do not modify the current worktree."""
66
67        if revision == "current":
68            print("Using current work directory.")
69            git_worktree_path = self.repo_path
70        else:
71            print("Creating git worktree for", revision)
72            git_worktree_path = os.path.join(self.repo_path, "temp-" + revision)
73            subprocess.check_output(
74                [self.git_command, "worktree", "add", "--detach",
75                 git_worktree_path, revision], cwd=self.repo_path,
76                stderr=subprocess.STDOUT
77            )
78        return git_worktree_path
79
80    def _build_libraries(self, git_worktree_path):
81        """Build libraries in the specified worktree."""
82
83        my_environment = os.environ.copy()
84        subprocess.check_output(
85            [self.make_command, "-j", "lib"], env=my_environment,
86            cwd=git_worktree_path, stderr=subprocess.STDOUT,
87        )
88
89    def _gen_code_size_csv(self, revision, git_worktree_path):
90        """Generate code size csv file."""
91
92        csv_fname = revision + ".csv"
93        if revision == "current":
94            print("Measuring code size in current work directory.")
95        else:
96            print("Measuring code size for", revision)
97        result = subprocess.check_output(
98            ["size library/*.o"], cwd=git_worktree_path, shell=True
99        )
100        size_text = result.decode()
101        csv_file = open(os.path.join(self.csv_dir, csv_fname), "w")
102        for line in size_text.splitlines()[1:]:
103            data = line.split()
104            csv_file.write("{}, {}\n".format(data[5], data[3]))
105
106    def _remove_worktree(self, git_worktree_path):
107        """Remove temporary worktree."""
108        if git_worktree_path != self.repo_path:
109            print("Removing temporary worktree", git_worktree_path)
110            subprocess.check_output(
111                [self.git_command, "worktree", "remove", "--force",
112                 git_worktree_path], cwd=self.repo_path,
113                stderr=subprocess.STDOUT
114            )
115
116    def _get_code_size_for_rev(self, revision):
117        """Generate code size csv file for the specified git revision."""
118
119        # Check if the corresponding record exists
120        csv_fname = revision + ".csv"
121        if (revision != "current") and \
122           os.path.exists(os.path.join(self.csv_dir, csv_fname)):
123            print("Code size csv file for", revision, "already exists.")
124        else:
125            git_worktree_path = self._create_git_worktree(revision)
126            self._build_libraries(git_worktree_path)
127            self._gen_code_size_csv(revision, git_worktree_path)
128            self._remove_worktree(git_worktree_path)
129
130    def compare_code_size(self):
131        """Generate results of the size changes between two revisions,
132        old and new. Measured code size results of these two revisions
133        must be available."""
134
135        old_file = open(os.path.join(self.csv_dir, self.old_rev + ".csv"), "r")
136        new_file = open(os.path.join(self.csv_dir, self.new_rev + ".csv"), "r")
137        res_file = open(os.path.join(self.result_dir, "compare-" + self.old_rev
138                                     + "-" + self.new_rev + ".csv"), "w")
139
140        res_file.write("file_name, this_size, old_size, change, change %\n")
141        print("Generating comparison results.")
142
143        old_ds = {}
144        for line in old_file.readlines()[1:]:
145            cols = line.split(", ")
146            fname = cols[0]
147            size = int(cols[1])
148            if size != 0:
149                old_ds[fname] = size
150
151        new_ds = {}
152        for line in new_file.readlines()[1:]:
153            cols = line.split(", ")
154            fname = cols[0]
155            size = int(cols[1])
156            new_ds[fname] = size
157
158        for fname in new_ds:
159            this_size = new_ds[fname]
160            if fname in old_ds:
161                old_size = old_ds[fname]
162                change = this_size - old_size
163                change_pct = change / old_size
164                res_file.write("{}, {}, {}, {}, {:.2%}\n".format(fname, \
165                               this_size, old_size, change, float(change_pct)))
166            else:
167                res_file.write("{}, {}\n".format(fname, this_size))
168        return 0
169
170    def get_comparision_results(self):
171        """Compare size of library/*.o between self.old_rev and self.new_rev,
172        and generate the result file."""
173        build_tree.check_repo_path()
174        self._get_code_size_for_rev(self.old_rev)
175        self._get_code_size_for_rev(self.new_rev)
176        return self.compare_code_size()
177
178def main():
179    parser = argparse.ArgumentParser(
180        description=(
181            """This script is for comparing the size of the library files
182            from two different Git revisions within an Mbed TLS repository.
183            The results of the comparison is formatted as csv, and stored at
184            a configurable location.
185            Note: must be run from Mbed TLS root."""
186        )
187    )
188    parser.add_argument(
189        "-r", "--result-dir", type=str, default="comparison",
190        help="directory where comparison result is stored, \
191              default is comparison",
192    )
193    parser.add_argument(
194        "-o", "--old-rev", type=str, help="old revision for comparison.",
195        required=True,
196    )
197    parser.add_argument(
198        "-n", "--new-rev", type=str, default=None,
199        help="new revision for comparison, default is the current work \
200              directory, including uncommitted changes."
201    )
202    comp_args = parser.parse_args()
203
204    if os.path.isfile(comp_args.result_dir):
205        print("Error: {} is not a directory".format(comp_args.result_dir))
206        parser.exit()
207
208    validate_res = CodeSizeComparison.validate_revision(comp_args.old_rev)
209    old_revision = validate_res.decode().replace("\n", "")
210
211    if comp_args.new_rev is not None:
212        validate_res = CodeSizeComparison.validate_revision(comp_args.new_rev)
213        new_revision = validate_res.decode().replace("\n", "")
214    else:
215        new_revision = "current"
216
217    result_dir = comp_args.result_dir
218    size_compare = CodeSizeComparison(old_revision, new_revision, result_dir)
219    return_code = size_compare.get_comparision_results()
220    sys.exit(return_code)
221
222
223if __name__ == "__main__":
224    main()
225