1 files changed, 36 insertions, 27 deletions
diff --git a/src/Haskell/AST.hs b/src/Haskell/AST.hs
index 2238b6d..6d25153 100644
--- a/src/Haskell/AST.hs
+++ b/src/Haskell/AST.hs
@@ -108,33 +108,42 @@ instance Pretty Inst where
     pretty (Inst n t ds) = Node ("instance " ++ n ++ " " ++ pprintOneline t ++ " where") [Bracket "{" "}" ";" (map pretty ds)]
 
 
-class AllRefs a where
-    allRefs :: a -> [Name]
-
-instance AllRefs AST where
-    allRefs (AST tops) = nub $ concatMap allRefs tops
-
-instance AllRefs Toplevel where
-    allRefs (TopDef def) = allRefs def
-    allRefs (TopDecl _) = []
-    allRefs (TopData _) = []
-    allRefs (TopClass _) = []
-    allRefs (TopInst inst) = allRefs inst
-
-instance AllRefs Def where
-    allRefs (Def _ e) = allRefs e
-
-instance AllRefs Expr where
-    allRefs (App e es) = nub $ concatMap allRefs (e : es)
-    allRefs (Ref n) = [n]
-    allVars (Con _) = []
-    allRefs (Num _) = []
-    allRefs (Tup es) = nub $ concatMap allRefs es
-    allRefs (Lam ns e) = allRefs e \\ ns
-    allRefs (Case e pairs) = nub $ allRefs e ++ concatMap (allRefs . snd) pairs
-
-instance AllRefs Inst where
-    allRefs (Inst _ _ ds) = nub $ concatMap allRefs ds
+-- This excludes constructor names, since those are not variables. This _does_
+-- include bound variables; if you don't want that, use freeVariables.
+class AllVars a where
+    allVars :: a -> Set.Set Name
+
+instance AllVars AST where
+    allVars (AST tops) = Set.unions (map allVars tops)
+
+instance AllVars Toplevel where
+    allVars (TopDef def) = allVars def
+    allVars (TopDecl _) = mempty
+    allVars (TopData _) = mempty
+    allVars (TopClass _) = mempty
+    allVars (TopInst inst) = allVars inst
+
+instance AllVars Def where
+    allVars (Def n e) = Set.insert n (allVars e)
+
+instance AllVars Inst where
+    allVars (Inst _ _ ds) = Set.unions (map allVars ds)
+
+instance AllVars Expr where
+    allVars (App e es) = Set.unions (map allVars (e : es))
+    allVars (Ref n) = Set.singleton n
+    allVars (Con _) = mempty
+    allVars (Num _) = mempty
+    allVars (Tup es) = Set.unions (map allVars es)
+    allVars (Lam ns e) = Set.fromList ns <> allVars e
+    allVars (Case e pairs) =
+        allVars e <> Set.unions [allVars p <> allVars e' | (p, e') <- pairs]
+
+instance AllVars Pat where
+    allVars PatAny = mempty
+    allVars (PatVar n) = Set.singleton n
+    allVars (PatCon _ ps) = Set.unions (map allVars ps)
+    allVars (PatTup ps) = Set.unions (map allVars ps)
 
 
 boundVars :: Pat -> Set.Set Name