From e23c4d1c63ff7fc3e14f655ccde91eec2c945edc Mon Sep 17 00:00:00 2001
From: Troels Henriksen <athas@sigkill.dk>
Date: Tue, 13 Jun 2023 23:33:43 +0200
Subject: [PATCH] Also look in expression.

---
 src/Language/Futhark/Query.hs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/Language/Futhark/Query.hs b/src/Language/Futhark/Query.hs
index 1fd048e5ae..7adb7b2686 100644
--- a/src/Language/Futhark/Query.hs
+++ b/src/Language/Futhark/Query.hs
@@ -230,7 +230,7 @@ atPosInTypeExp te pos =
       msum $ map (`atPosInTypeExp` pos) $ concatMap snd cs
     TEDim _ t _ ->
       atPosInTypeExp t pos
-    TERefine t _ _ -> atPosInTypeExp t pos
+    TERefine t e _ -> atPosInTypeExp t pos `mplus` atPosInExp e pos
   where
     inArg (TypeArgExpSize dim) = inDim dim
     inArg (TypeArgExpType e2) = atPosInTypeExp e2 pos