decomp: Support LZCOUNT & POPCOUNT ops (#118)

* decomp: Add builtin call helper

- Supports adding arbitrary builtins, fixing an issue with trying to add LZCOUNT & POPCOUNT support

* decomp: Support LZCOUNT & POPCOUNT

- Relies on emission of clang's builtin functions for the pcode op

Author: kyle-elliott-tob

include/patchestry/AST/OperationBuilder.hpp

@@ -156,6 +156,12 @@ namespace patchestry::ast {
             clang::ASTContext &ctx, const Function &function, const Operation &op
         );
 
+        std::pair< clang::Stmt *, bool >
+        create_popcount(clang::ASTContext &ctx, const Function &function, const Operation &op);
+
+        std::pair< clang::Stmt *, bool >
+        create_lzcount(clang::ASTContext &ctx, const Function &function, const Operation &op);
+
         std::pair< clang::Stmt *, bool > create_builtin_call_expr(
             clang::ASTContext &ctx, const Function &function, const Operation &op,
             clang::Builtin::ID id

include/patchestry/Ghidra/Pcode.def

@@ -76,4 +76,6 @@
         X(DECLARE_PARAMETER) \
         X(DECLARE_LOCAL) \
         X(DECLARE_TEMPORARY) \
-        X(ADDRESS_OF)
+        X(ADDRESS_OF) \
+        X(POPCOUNT) \
+        X(LZCOUNT)

lib/patchestry/AST/FunctionBuilder.cpp

@@ -643,6 +643,10 @@ namespace patchestry::ast {
                 return op_builder->create_declare_local(ctx, op);
             case Mnemonic::OP_ADDRESS_OF:
                 return op_builder->create_unary_operation(ctx, function, op, clang::UO_AddrOf);
+            case Mnemonic::OP_POPCOUNT:
+                return op_builder->create_popcount(ctx, function, op);
+            case Mnemonic::OP_LZCOUNT:
+                return op_builder->create_lzcount(ctx, function, op);
             case Mnemonic::OP_UNKNOWN:
                 assert(false);
                 break;

lib/patchestry/AST/OperationStmt.cpp

@@ -22,6 +22,7 @@
 #include <clang/Basic/LangOptions.h>
 #include <clang/Basic/SourceLocation.h>
 #include <clang/Basic/Specifiers.h>
+#include <clang/Sema/Lookup.h>
 #include <clang/Sema/Sema.h>
 #include <llvm/ADT/APInt.h>
 #include <llvm/Support/Casting.h>
@@ -89,6 +90,42 @@ namespace patchestry::ast {
             }
         }
 
+        /**
+         * Helper function to create a builtin call expression.
+         */
+        clang::Expr *create_builtin_call(
+            const clang::ASTContext &ctx, clang::Sema &sema,
+            const clang::Builtin::ID builtin_id, std::vector< clang::Expr * > &args,
+            const clang::SourceLocation loc
+        ) {
+            // Get the builtin name and create a lookup result
+            const auto name = ctx.BuiltinInfo.getName(builtin_id);
+            const clang::LookupResult r(
+                sema, &ctx.Idents.get(name), loc, clang::Sema::LookupOrdinaryName
+            );
+            auto *II = r.getLookupName().getAsIdentifierInfo();
+
+            // Get the builtin type
+            clang::ASTContext::GetBuiltinTypeError error{};
+            const auto ty = ctx.GetBuiltinType(builtin_id, error);
+            assert(!error && "Failed to get builtin type");
+
+            // Create the builtin declaration
+            auto *builtin_decl = sema.CreateBuiltin(II, ty, builtin_id, loc);
+
+            // Create a DeclRefExpr for the builtin
+            auto *decl_ref = clang::DeclRefExpr::Create(
+                ctx, clang::NestedNameSpecifierLoc(), clang::SourceLocation(), builtin_decl,
+                false, loc, builtin_decl->getType(), clang::VK_PRValue
+            );
+
+            // Build the call expression
+            auto result = sema.BuildCallExpr(nullptr, decl_ref, loc, args, loc);
+            assert(!result.isInvalid() && "Failed to build builtin call expr");
+
+            return result.getAs< clang::Expr >();
+        }
+
     } // namespace
 
     /**
@@ -1215,6 +1252,34 @@ namespace patchestry::ast {
         return create_builtin_call_expr(ctx, function, op, clang::Builtin::BIroundf);
     }
 
+    std::pair< clang::Stmt *, bool > OpBuilder::create_popcount(
+        clang::ASTContext &ctx, const Function &function, const Operation &op
+    ) {
+        if (op.inputs.size() != 1U || op.mnemonic != Mnemonic::OP_POPCOUNT) {
+            LOG(ERROR) << "POPCOUNT operation is invalid or has invalid input operand. key: "
+                       << op.key << "\n";
+            return {};
+        }
+
+        // Use __builtin_popcount for population count
+        return create_builtin_call_expr(
+            ctx, function, op, clang::Builtin::BI__builtin_popcount
+        );
+    }
+
+    std::pair< clang::Stmt *, bool > OpBuilder::create_lzcount(
+        clang::ASTContext &ctx, const Function &function, const Operation &op
+    ) {
+        if (op.inputs.size() != 1U || op.mnemonic != Mnemonic::OP_LZCOUNT) {
+            LOG(ERROR) << "LZCOUNT operation is invalid or has invalid input operand. key: "
+                       << op.key << "\n";
+            return {};
+        }
+
+        // Use __builtin_clz for count leading zeros
+        return create_builtin_call_expr(ctx, function, op, clang::Builtin::BI__builtin_clz);
+    }
+
     std::pair< clang::Stmt *, bool > OpBuilder::create_builtin_call_expr(
         clang::ASTContext &ctx, const Function &function, const Operation &op,
         clang::Builtin::ID id
@@ -1224,7 +1289,9 @@ namespace patchestry::ast {
         auto *input_expr =
             clang::dyn_cast< clang::Expr >(create_varnode(ctx, function, op.inputs[0]));
 
-        auto *call_expr = sema().BuildBuiltinCallExpr(op_loc, id, { input_expr });
+        std::vector args = { input_expr };
+        auto *call_expr = create_builtin_call(ctx, sema(), id, args, op_loc);
+
         if (op.output.has_value()) {
             return { call_expr, true };
         }