Fix issues with extract() deleting used arguments

Author: devajithvs

mlir/lib/Tools/mlir-query/Marshallers.h

@@ -144,7 +144,7 @@ class VariadicMatcherCreateCallback : public MatcherCreateCallback {
 
 // Helper function to perform template argument deduction.
 template <typename MarshallerType, typename FuncType>
-MatcherCreateCallback *createMarshallerCallback(MarshallerType Marshaller,
+auto createMarshallerCallback(MarshallerType Marshaller,
                                                 FuncType Func,
                                                 StringRef MatcherName) {
   return new FixedArgCountMatcherCreateCallback<MarshallerType, FuncType>(
@@ -235,23 +235,23 @@ DynMatcher *matcherMarshall1(ReturnType (*Func)(InArgType1, InArgType2),
 
 // 0-arg overload
 template <typename T, typename ReturnType>
-MatcherCreateCallback *makeMatcherAutoMarshall(ReturnType (*Func)(),
+auto makeMatcherAutoMarshall(ReturnType (*Func)(),
                                                StringRef MatcherName) {
   return createMarshallerCallback(matcherMarshall0<T, ReturnType>, Func,
                                   MatcherName);
 }
 
 // 1-arg overload
 template <typename T, typename ReturnType, typename ArgType1>
-MatcherCreateCallback *makeMatcherAutoMarshall(ReturnType (*Func)(ArgType1),
+auto makeMatcherAutoMarshall(ReturnType (*Func)(ArgType1),
                                                StringRef MatcherName) {
   return createMarshallerCallback(matcherMarshall1<T, ReturnType, ArgType1>,
                                   Func, MatcherName);
 }
 
 // 2-arg overload
 template <typename T, typename ReturnType, typename ArgType1, typename ArgType2>
-MatcherCreateCallback *makeMatcherAutoMarshall(ReturnType (*Func)(ArgType1,
+auto makeMatcherAutoMarshall(ReturnType (*Func)(ArgType1,
                                                                   ArgType2),
                                                StringRef MatcherName) {
   return createMarshallerCallback(
@@ -260,7 +260,7 @@ MatcherCreateCallback *makeMatcherAutoMarshall(ReturnType (*Func)(ArgType1,
 
 // Variadic overload.
 template <typename T, typename MatcherType>
-MatcherCreateCallback *makeMatcherAutoMarshall(MatcherType Func,
+auto makeMatcherAutoMarshall(MatcherType Func,
                                                StringRef MatcherName) {
   return new VariadicMatcherCreateCallback<T>(MatcherName);
 }

mlir/lib/Tools/mlir-query/Query.cpp

@@ -95,9 +95,9 @@ Operation *extractFunction(std::vector<matcher::DynTypedNode> &nodes,
     builder.clone(*slicedOp, mapper);
 
   // Remove func arguments that are not used.
-  int currentIndex = 0;
-  for (auto value : funcOp.getArguments()) {
-    if (value.getUses().empty()) {
+  unsigned currentIndex = 0;
+  while (currentIndex < funcOp.getNumArguments()) {
+    if (funcOp.getArgument(currentIndex).getUses().empty()) {
       funcOp.eraseArgument(currentIndex);
     } else {
       currentIndex++;