1#!/usr/bin/env python3 2# -- coding: utf-8 -- 3# 4# Copyright (c) 2025 Huawei Device Co., Ltd. 5# Licensed under the Apache License, Version 2.0 (the "License"); 6# you may not use this file except in compliance with the License. 7# You may obtain a copy of the License at 8# 9# http://www.apache.org/licenses/LICENSE-2.0 10# 11# Unless required by applicable law or agreed to in writing, software 12# distributed under the License is distributed on an "AS IS" BASIS, 13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14# See the License for the specific language governing permissions and 15# limitations under the License. 16# 17 18import sys 19import os 20import uuid 21from pathlib import Path 22from typing import Optional, List 23from multiprocessing import current_process 24 25 26BINARY_NAMES: set[str] = {'es2panda', 'ark', 'verifier', 'ark_aot', 'ark_asm'} 27SCRIPT_FILE_EXTENSIONS: set[str] = {'.rb', '.py', '.sh'} 28COVERAGE_ROOT_DIR: Path = Path('/tmp/coverage_intermediate') 29LLVM_PROFDATA_BINARY: str = 'llvm-profdata-14' 30 31COVERAGE_TOOLS = {'llvm-cov', 'lcov'} 32 33ADD_CUSTOM_TARGET_FUNCTION_ARGS: set[str] = { 34 'ALL', 'DEPENDS', 'BYPRODUCTS', 'WORKING_DIRECTORY', 35 'COMMENT', 'VERBATIM', 'USES_TERMINAL', 36 'COMMAND_EXPAND_LISTS', 'JOB_SERVER_AWARE', 'COMMAND' 37} 38 39 40def is_script_file(file_path: str) -> bool: 41 return any(file_path.endswith(ext) for ext in SCRIPT_FILE_EXTENSIONS) 42 43 44def is_script_call(command_part: str) -> bool: 45 return (is_script_file(command_part) or 46 (command_part.startswith('./') and is_script_file(command_part))) 47 48 49def is_cmake_copy_command(command_parts: List[str]) -> bool: 50 return (len(command_parts) >= 4 and 51 command_parts[0].endswith('cmake') and 52 command_parts[1] == '-E' and 53 command_parts[2] in {'copy', 'copy_directory'}) 54 55 56def is_target_binary(command_part: str) -> bool: 57 """Check if the command part refers to one of our target binaries.""" 58 command_part = command_part.lower() 59 60 # Direct binary name match 61 if command_part in {cmd.lower() for cmd in BINARY_NAMES}: 62 return True 63 64 # Check for CMake generator expressions 65 if '$<' in command_part and '>' in command_part: 66 return any(cmd.lower() in command_part for cmd in BINARY_NAMES) 67 68 # Check for binary in file path 69 if '/' in command_part or '\\' in command_part: 70 filename = os.path.basename(command_part).lower() 71 return filename in {cmd.lower() for cmd in BINARY_NAMES} 72 73 return False 74 75 76def generate_profdata_merge_command(profraw_path: Path) -> str: 77 """Generate command to merge LLVM profraw files into profdata format.""" 78 return (f"COMMAND;{LLVM_PROFDATA_BINARY};merge;" 79 f"-output={profraw_path.with_suffix('.profdata')};{profraw_path}") 80 81 82def extract_target_binary_name(command_part: str) -> Optional[str]: 83 """Extract the binary name from a command part if it's one of our targets.""" 84 85 # Handle CMake generator expressions 86 if command_part.startswith('$<TARGET_FILE:') and command_part.endswith('>'): 87 binary_name = command_part[len('$<TARGET_FILE:'):-1] 88 if binary_name in BINARY_NAMES: 89 return binary_name 90 91 # Direct binary name match 92 if command_part in BINARY_NAMES: 93 return command_part 94 95 # Extract from file path 96 if '/' in command_part or '\\' in command_part: 97 filename = os.path.basename(command_part) 98 return next((cmd for cmd in BINARY_NAMES 99 if cmd.lower() == filename.lower()), filename) 100 101 return command_part 102 103 104def find_binary_in_command(command_parts: List[str]) -> Optional[str]: 105 for part in command_parts: 106 binary_name = extract_target_binary_name(part) 107 if binary_name is not None and binary_name in BINARY_NAMES: 108 return binary_name 109 return None 110 111 112def should_skip_command_processing(command_parts: List[str]) -> bool: 113 return is_cmake_copy_command(command_parts) or any( 114 is_script_call(part) for part in command_parts) 115 116 117def instrument_for_lcov_coverage(command_parts: List[str]) -> List[str]: 118 if should_skip_command_processing(command_parts): 119 return command_parts 120 121 binary_name = find_binary_in_command(command_parts) 122 if binary_name is None: 123 return command_parts 124 125 gcov_dir_path = COVERAGE_ROOT_DIR / Path(binary_name) 126 modified_parts = command_parts.copy() 127 for i, part in enumerate(modified_parts): 128 if part.startswith('LD_LIBRARY_PATH=') or is_target_binary(part): 129 modified_parts.insert(i, f"GCOV_PREFIX={gcov_dir_path}") 130 break 131 return modified_parts 132 133 134def instrument_for_llvm_coverage(command_parts: List[str]) -> List[str]: 135 if should_skip_command_processing(command_parts): 136 return command_parts 137 138 binary_name = find_binary_in_command(command_parts) 139 if binary_name is None: 140 return command_parts 141 142 pid = current_process().pid 143 unique_id = uuid.uuid4() 144 file_path = COVERAGE_ROOT_DIR / Path(f"{binary_name}-{pid}-{unique_id}") 145 profraw_path = file_path.with_suffix('.profraw') 146 147 modified_parts = command_parts.copy() 148 for i, part in enumerate(modified_parts): 149 if part.startswith('LD_LIBRARY_PATH=') or is_target_binary(part): 150 modified_parts.insert(i, f"LLVM_PROFILE_FILE={profraw_path}") 151 modified_parts.extend([ 152 generate_profdata_merge_command(profraw_path), 153 f"COMMAND;rm;{profraw_path}" 154 ]) 155 break 156 return modified_parts 157 158 159def process_command_section(command_parts: List[str], tool: str) -> List[str]: 160 """Process COMMAND section based on the selected coverage tool.""" 161 processors = { 162 'llvm-cov': instrument_for_llvm_coverage, 163 'lcov': instrument_for_lcov_coverage 164 } 165 processor = processors.get(tool) 166 if processor is None: 167 raise ValueError(f"Unsupported coverage tool: {tool}. " 168 f"Supported tools are {list(processors.keys())}") 169 return processor(command_parts) 170 171 172def process_input_sections(input_string: str, tool: str) -> List[str]: 173 sections = input_string.strip().split(';') 174 processed_sections = [] 175 i = 0 176 while i < len(sections): 177 section = sections[i] 178 if section not in ADD_CUSTOM_TARGET_FUNCTION_ARGS: 179 processed_sections.append(section) 180 i += 1 181 continue 182 if section == 'COMMAND': 183 command_parts = [] 184 i += 1 185 while i < len(sections) and sections[i] not in ADD_CUSTOM_TARGET_FUNCTION_ARGS: 186 command_parts.append(sections[i]) 187 i += 1 188 189 processed_command = process_command_section(command_parts, tool) 190 processed_sections.extend(['COMMAND'] + processed_command) 191 else: 192 end = i + 1 193 while end < len(sections) and sections[end] not in ADD_CUSTOM_TARGET_FUNCTION_ARGS: 194 end += 1 195 processed_sections.extend(sections[i:end]) 196 i = end 197 return processed_sections 198 199 200def main(tool: str, input_string: str) -> None: 201 if tool not in COVERAGE_TOOLS: 202 raise ValueError("Tool must be either 'llvm-cov' or 'lcov'") 203 processed_sections = process_input_sections(input_string, tool) 204 print(';'.join(processed_sections)) 205 206 207if __name__ == '__main__': 208 if len(sys.argv) != 3: 209 print("Usage: python script.py 'tool'(llvm-cov/lcov) 'input_string'") 210 sys.exit(1) 211 try: 212 main(sys.argv[1], sys.argv[2]) 213 except ValueError as e: 214 print(f"Error: {e}") 215 sys.exit(1) 216