@@ -277,6 +277,7 @@ def _generate_implementation(self, analysis_result: AnalysisResult) -> str:
277277 impl = """#include "generated.hpp"
278278#include <vector>
279279#include <map>
280+ #include <unordered_map>
280281#include <set>
281282#include <tuple>
282283#include <optional>
@@ -1083,6 +1084,10 @@ def _translate_expression(self, node: ast.AST, local_vars: Dict[str, str]) -> st
10831084 obj = self ._translate_expression (node .func .value .value , local_vars )
10841085 args = [self ._translate_expression (arg , local_vars ) for arg in node .args ]
10851086 return f"{ obj } .push_back({ ', ' .join (args )} )"
1087+ elif func_name in ['sqrt' , 'sin' , 'cos' , 'tan' , 'asin' , 'acos' , 'atan' , 'exp' , 'log' , 'log10' , 'floor' , 'ceil' , 'fabs' ]:
1088+ # Handle direct imports from math module (e.g., from math import sqrt)
1089+ args = [self ._translate_expression (arg , local_vars ) for arg in node .args ]
1090+ return f"std::{ func_name } ({ ', ' .join (args )} )"
10861091 else :
10871092 # Regular function call
10881093 args = [self ._translate_expression (arg , local_vars ) for arg in node .args ]
@@ -1092,6 +1097,16 @@ def _translate_expression(self, node: ast.AST, local_vars: Dict[str, str]) -> st
10921097 obj = self ._translate_expression (node .func .value , local_vars )
10931098 method = node .func .attr
10941099
1100+ # Handle math module functions
1101+ if obj == 'math' :
1102+ if method in ['sqrt' , 'sin' , 'cos' , 'tan' , 'asin' , 'acos' , 'atan' , 'exp' , 'log' , 'log10' , 'pow' , 'floor' , 'ceil' , 'fabs' ]:
1103+ args = [self ._translate_expression (arg , local_vars ) for arg in node .args ]
1104+ return f"std::{ method } ({ ', ' .join (args )} )"
1105+ else :
1106+ # Handle other math functions that may need special mapping
1107+ args = [self ._translate_expression (arg , local_vars ) for arg in node .args ]
1108+ return f"std::{ method } ({ ', ' .join (args )} )"
1109+
10951110 # Map Python methods to C++ equivalents
10961111 if method == 'append' :
10971112 method = 'push_back' # std::vector uses push_back, not append
@@ -1153,6 +1168,12 @@ def _translate_expression(self, node: ast.AST, local_vars: Dict[str, str]) -> st
11531168 return "std::make_tuple()"
11541169
11551170 return f"std::make_tuple({ ', ' .join (elements )} )"
1171+ elif isinstance (node , ast .ListComp ):
1172+ # Handle list comprehensions: [expr for item in iterable]
1173+ return self ._translate_list_comprehension (node , local_vars )
1174+ elif isinstance (node , ast .DictComp ):
1175+ # Handle dictionary comprehensions: {key: value for item in iterable}
1176+ return self ._translate_dict_comprehension (node , local_vars )
11561177 elif isinstance (node , ast .BoolOp ):
11571178 # Handle boolean operations like and, or
11581179 op_str = "&&" if isinstance (node .op , ast .And ) else "||"
@@ -1619,6 +1640,122 @@ def _generate_cmake(self) -> str:
16191640
16201641 return '\n ' .join (cmake_content )
16211642
1643+ def _translate_list_comprehension (self , node : ast .ListComp , local_vars : Dict [str , str ]) -> str :
1644+ """Translate list comprehension to C++ lambda with performance optimizations."""
1645+ # Get the comprehension parts
1646+ element_expr = node .elt
1647+ generator = node .generators [0 ] # For simplicity, handle only one generator
1648+ target = generator .target
1649+ iter_expr = generator .iter
1650+
1651+ # Translate components
1652+ iter_str = self ._translate_expression (iter_expr , local_vars )
1653+ target_name = target .id if isinstance (target , ast .Name ) else "item"
1654+ element_str = self ._translate_expression (element_expr , local_vars )
1655+
1656+ # Handle conditional comprehensions (if clauses)
1657+ condition_str = ""
1658+ if generator .ifs :
1659+ conditions = []
1660+ for if_clause in generator .ifs :
1661+ condition = self ._translate_expression (if_clause , local_vars )
1662+ conditions .append (condition )
1663+ condition_str = f" if ({ ' && ' .join (conditions )} )"
1664+
1665+ # Create lambda expression for the comprehension with performance optimizations
1666+ # [expr for item in iterable] becomes:
1667+ # [&]() {
1668+ # std::vector<auto> result;
1669+ # result.reserve(iterable.size()); // Performance optimization
1670+ # for (auto item : iterable) {
1671+ # if (condition) { // Only if conditions exist
1672+ # result.push_back(expr);
1673+ # }
1674+ # }
1675+ # return result;
1676+ # }()
1677+
1678+ if condition_str :
1679+ comprehension_code = f"""[&]() {{
1680+ std::vector<auto> result;
1681+ result.reserve({ iter_str } .size());
1682+ for (auto { target_name } : { iter_str } ) {{
1683+ if ({ ' && ' .join (self ._translate_expression (if_clause , local_vars ) for if_clause in generator .ifs )} ) {{
1684+ result.push_back({ element_str } );
1685+ }}
1686+ }}
1687+ return result;
1688+ }}()"""
1689+ else :
1690+ comprehension_code = f"""[&]() {{
1691+ std::vector<auto> result;
1692+ result.reserve({ iter_str } .size());
1693+ for (auto { target_name } : { iter_str } ) {{
1694+ result.push_back({ element_str } );
1695+ }}
1696+ return result;
1697+ }}()"""
1698+
1699+ return comprehension_code
1700+
1701+ def _translate_dict_comprehension (self , node : ast .DictComp , local_vars : Dict [str , str ]) -> str :
1702+ """Translate dictionary comprehension to C++ lambda with performance optimizations."""
1703+ # Get the comprehension parts
1704+ key_expr = node .key
1705+ value_expr = node .value
1706+ generator = node .generators [0 ] # For simplicity, handle only one generator
1707+ target = generator .target
1708+ iter_expr = generator .iter
1709+
1710+ # Translate components
1711+ iter_str = self ._translate_expression (iter_expr , local_vars )
1712+ target_name = target .id if isinstance (target , ast .Name ) else "item"
1713+ key_str = self ._translate_expression (key_expr , local_vars )
1714+ value_str = self ._translate_expression (value_expr , local_vars )
1715+
1716+ # Handle conditional comprehensions (if clauses)
1717+ condition_str = ""
1718+ if generator .ifs :
1719+ conditions = []
1720+ for if_clause in generator .ifs :
1721+ condition = self ._translate_expression (if_clause , local_vars )
1722+ conditions .append (condition )
1723+ condition_str = f" if ({ ' && ' .join (conditions )} )"
1724+
1725+ # Create lambda expression for the dictionary comprehension with performance optimizations
1726+ # Use std::unordered_map instead of std::map for O(1) vs O(log n) performance
1727+ # {key: value for item in iterable} becomes:
1728+ # [&]() {
1729+ # std::unordered_map<auto, auto> result;
1730+ # for (auto item : iterable) {
1731+ # if (condition) { // Only if conditions exist
1732+ # result[key] = value;
1733+ # }
1734+ # }
1735+ # return result;
1736+ # }()
1737+
1738+ if condition_str :
1739+ comprehension_code = f"""[&]() {{
1740+ std::unordered_map<auto, auto> result;
1741+ for (auto { target_name } : { iter_str } ) {{
1742+ if ({ ' && ' .join (self ._translate_expression (if_clause , local_vars ) for if_clause in generator .ifs )} ) {{
1743+ result[{ key_str } ] = { value_str } ;
1744+ }}
1745+ }}
1746+ return result;
1747+ }}()"""
1748+ else :
1749+ comprehension_code = f"""[&]() {{
1750+ std::unordered_map<auto, auto> result;
1751+ for (auto { target_name } : { iter_str } ) {{
1752+ result[{ key_str } ] = { value_str } ;
1753+ }}
1754+ return result;
1755+ }}()"""
1756+
1757+ return comprehension_code
1758+
16221759 def _expression_uses_variables (self , expr : ast .AST , variable_names : List [str ]) -> bool :
16231760 """Check if an expression uses any of the given variable names."""
16241761 for node in ast .walk (expr ):
0 commit comments