Merge pull request #1324 from ZenithalHourlyRate:follow-1306

PiperOrigin-RevId: 720545594

Author: copybara-github

lib/Dialect/LWE/Conversions/LWEToOpenfhe/LWEToOpenfhe.cpp

@@ -92,23 +92,14 @@ struct AddCryptoContextArg : public OpConversionPattern<func::FuncOp> {
     }
 
     auto cryptoContextType = openfhe::CryptoContextType::get(getContext());
-    FunctionType originalType = op.getFunctionType();
-    llvm::SmallVector<Type, 4> newTypes;
-    newTypes.reserve(originalType.getNumInputs() + 1);
-    newTypes.push_back(cryptoContextType);
-    for (auto t : originalType.getInputs()) {
-      newTypes.push_back(t);
-    }
-    auto newFuncType =
-        FunctionType::get(getContext(), newTypes, originalType.getResults());
     rewriter.modifyOpInPlace(op, [&] {
-      op.setType(newFuncType);
-
-      // guard against private FuncOp (i.e. declaration)
-      if (op.getVisibility() != SymbolTable::Visibility::Private) {
-        Block &block = op.getBody().getBlocks().front();
-        block.insertArgument(&block.getArguments().front(), cryptoContextType,
-                             op.getLoc());
+      if (op.isDeclaration()) {
+        auto newFuncType = op.getTypeWithArgsAndResults(
+            ArrayRef<unsigned int>{0}, ArrayRef<Type>{cryptoContextType}, {},
+            {});
+        op.setType(newFuncType);
+      } else {
+        op.insertArgument(0, cryptoContextType, nullptr, op.getLoc());
       }
     });
 

lib/Target/OpenFhePke/OpenFhePkeEmitter.cpp

@@ -171,7 +171,7 @@ LogicalResult OpenFhePkeEmitter::printOperation(func::FuncOp funcOp) {
     }
   }
 
-  if (funcOp.getVisibility() == SymbolTable::Visibility::Private) {
+  if (funcOp.isDeclaration()) {
     // function declaration
     os << commaSeparatedTypes(funcOp.getArgumentTypes(), [&](Type type) {
       return convertType(type, funcOp->getLoc()).value();
@@ -186,7 +186,7 @@ LogicalResult OpenFhePkeEmitter::printOperation(func::FuncOp funcOp) {
   os << ")";
 
   // function declaration
-  if (funcOp.getVisibility() == SymbolTable::Visibility::Private) {
+  if (funcOp.isDeclaration()) {
     os << ";\n";
     return success();
   }
@@ -213,7 +213,7 @@ LogicalResult OpenFhePkeEmitter::printOperation(func::CallOp op) {
   }
 
   if (op.getNumResults() != 0) {
-    os << variableNames->getNameForValue(op.getResult(0)) << " = ";
+    emitAutoAssignPrefix(op.getResult(0));
   }
 
   os << canonicalizeDebugPort(op.getCallee()) << "(";