More root pathing fixups

Author: wsmoses

src/compiler.jl

@@ -3802,58 +3802,59 @@ function move_sret_tofrom_roots!(builder::LLVM.IRBuilder, jltype::LLVM.LLVMType,
 	# aka bfs/etc
         while length(todo) != 0
             path, ty = popfirst!(todo)
+            if !any_jltypes(ty)
+                continue
+            end
+
             if isa(ty, LLVM.PointerType)
-		if direction == SRetPointerToRootPointer || direction == SRetValueToRootPointer || direction == RootPointerToSRetPointer || direction == RootPointerToSRetValue
-                  loc = inbounds_gep!(
-                      builder,
-                      root_ty,
-                      rootRet,
-		      to_llvm(Cuint[count]),
-		     )
-		end
-                
-		if direction == SRetPointerToRootPointer
-		    outloc = inbounds_gep!(builder, jltype, sret, to_llvm(path))
-		    outloc = load!(builder, ty, outloc)
-                    store!(builder, outloc, loc)
-		elseif direction == SRetValueToRootPointer
-		    outloc = Enzyme.API.e_extract_value!(builder, sret, path)
-                    store!(builder, outloc, loc)
-		elseif direction == RootPointerToSRetValue
-		    loc = load!(builder, ty, loc)
-		    val = Enzyme.API.e_insert_value!(builder, val, loc, path)
-        elseif direction == NullifySRetValue
-            loc = unsafe_to_llvm(builder, nothing)
-            val = Enzyme.API.e_insert_value!(builder, val, loc, path)
-		elseif direction == RootPointerToSRetPointer
-		    outloc = inbounds_gep!(builder, jltype, sret, to_llvm(path))
-		    loc = load!(builder, ty, loc)
-		    push!(extracted, loc)
-                    store!(builder, loc, outloc)
-		else
-		    @assert false "Unhandled direction"
-		end
-                
-		count += 1
+
+        		if direction == SRetPointerToRootPointer || direction == SRetValueToRootPointer || direction == RootPointerToSRetPointer || direction == RootPointerToSRetValue
+                          loc = inbounds_gep!(
+                              builder,
+                              root_ty,
+                              rootRet,
+        		      to_llvm(Cuint[count]),
+        		     )
+        		end
+                        
+        		if direction == SRetPointerToRootPointer
+        		    outloc = inbounds_gep!(builder, jltype, sret, to_llvm(path))
+        		    outloc = load!(builder, ty, outloc)
+                            store!(builder, outloc, loc)
+        		elseif direction == SRetValueToRootPointer
+        		    outloc = Enzyme.API.e_extract_value!(builder, sret, path)
+                            store!(builder, outloc, loc)
+        		elseif direction == RootPointerToSRetValue
+        		    loc = load!(builder, ty, loc)
+        		    val = Enzyme.API.e_insert_value!(builder, val, loc, path)
+                elseif direction == NullifySRetValue
+                    loc = unsafe_to_llvm(builder, nothing)
+                    val = Enzyme.API.e_insert_value!(builder, val, loc, path)
+        		elseif direction == RootPointerToSRetPointer
+        		    outloc = inbounds_gep!(builder, jltype, sret, to_llvm(path))
+        		    loc = load!(builder, ty, loc)
+        		    push!(extracted, loc)
+                            store!(builder, loc, outloc)
+        		else
+        		    @assert false "Unhandled direction"
+        		end
+                        
+        		count += 1
                 continue
             end
             if isa(ty, LLVM.ArrayType)
-                if any_jltypes(ty)
-                    for i = 1:length(ty)
-                        npath = copy(path)
-			push!(npath, i - 1)
-                        push!(todo, (npath, eltype(ty)))
-                    end
+                for i = 1:length(ty)
+                    npath = copy(path)
+		push!(npath, i - 1)
+                    push!(todo, (npath, eltype(ty)))
                 end
                 continue
             end
             if isa(ty, LLVM.VectorType)
-                if any_jltypes(ty)
-                    for i = 1:size(ty)
-                        npath = copy(path)
-			push!(npath, i - 1)
-                        push!(todo, (npath, eltype(ty)))
-                    end
+                for i = 1:size(ty)
+                    npath = copy(path)
+		push!(npath, i - 1)
+                    push!(todo, (npath, eltype(ty)))
                 end
                 continue
             end

src/rules/customrules.jl

@@ -2051,7 +2051,7 @@ end
     fop = called_operand(orig)::LLVM.Function
     for (i, v) in enumerate(operands(orig)[1:end-1])
         if v == val
-            if !has_arg_attr(fop, i, StringAttribute("enzymejl_returnRoots"))
+            if true || !has_arg_attr(fop, i, StringAttribute("enzymejl_returnRoots"))
                 non_rooting_use = true
                 break
             end