diff --git a/lib/ast/astDef.ml b/lib/ast/astDef.ml
index 721e510..65f22d6 100644
--- a/lib/ast/astDef.ml
+++ b/lib/ast/astDef.ml
@@ -265,7 +265,6 @@ module Type = struct
       | Bot
       | Any
       | Var of QualIdent.t
-      | Set
       | Map
       | Fld
       | Data of QualIdent.t * variant_decl list
@@ -316,7 +315,6 @@ module Type = struct
     | Bot -> bot_type_string
     | Any -> any_type_string
     | Ref -> ref_type_string
-    | Set -> set_type_string
     | Map -> map_type_string
     | Fld -> fld_type_string
     | Perm -> perm_type_string
@@ -327,7 +325,7 @@ module Type = struct
 
   let rec pr_constr ppf t =
     match t with
-    | Int | Real | Num | Bool | Any | Bot | Ref | Perm | Var _ | Set | AtomicToken
+    | Int | Real | Num | Bool | Any | Bot | Ref | Perm | Var _ | AtomicToken
     | Map | Fld | Prod ->
         Stdlib.Format.fprintf ppf "%s" (to_name t)
     | Data (id, decls) ->
@@ -338,6 +336,8 @@ module Type = struct
   and pr ppf t =
     match t with
     | App (t1, [], _) -> pr_constr ppf t1
+    | App (Map, [t1; App (Bool, _, _)], _) ->
+      Stdlib.Format.fprintf ppf "Set[%a]" pr t1
     | App (Prod, ts, _) ->
       Stdlib.Format.fprintf ppf "(@[%a@])" (Print.pr_list_comma pr) ts
     | App (t1, ts, _) ->
@@ -388,7 +388,7 @@ module Type = struct
   let mk_any loc = App (Any, [], mk_attr loc)
   let mk_bot loc = App (Bot, [], mk_attr loc)
   let mk_ref loc = App (Ref, [], mk_attr loc)
-  let mk_set loc tp = App (Set, [tp], mk_attr loc)
+  let mk_set loc tp = App (Map, [tp; mk_bool loc], mk_attr loc)
   let mk_map loc tpi tpo = App (Map, [tpi; tpo], mk_attr loc)
   let mk_fld loc tpf = App (Fld, [tpf], mk_attr loc)
   let mk_perm loc = App (Perm, [], mk_attr loc)
@@ -444,7 +444,8 @@ module Type = struct
     match (t1, t2) with
     | App (Bot, [], _), t | t, App (Bot, [], _) -> t
     | App (t1, [], a1), App (t2, [], _) -> App (join_constr t1 t2, [], a1)
-    | App (Set, [t1], a1), App (Set, [t2], _a2) -> App (Set, [join t1 t2], a1)
+    | App (Map, [t1; App (Bool, [], a0)], a1), App (Map, [t2; App (Bool, [], _)], _a2) ->
+      App (Map, [join t1 t2; App (Bool, [], a0)], a1)
     | App (Map, [ti1; to1], a1), App (Map, [ti2; to2], _) -> App (Map, [meet ti1 ti2; join to1 to2], a1)
     | App (Prod, ts1, a1), App (Prod, ts2, _a2) ->
       (List.map2 ~f:join ts1 ts2 |> function
@@ -458,7 +459,8 @@ module Type = struct
     match (t1, t2) with
     | App (Any, [], _), t | t, App (Any, [], _) -> t
     | App (t1, [], a1), App (t2, [], _) -> App (meet_constr t1 t2, [], a1)
-    | App (Set, [t1], a1), App (Set, [t2], _a2) -> App (Set, [meet t1 t2], a1)
+    | App (Map, [t1; App (Bool, [], a0)], a1), App (Map, [t2; App (Bool, [], _)], _a2) ->
+      App (Map, [meet t1 t2; App (Bool, [], a0)], a1)
     | App (Map, [ti1; to1], a1), App (Map, [ti2; to2], _) -> App (Map, [join ti1 ti2; meet to1 to2], a1)
     | App (Prod, ts1, a1), App (Prod, ts2, _a2) ->
       (List.map2 ~f:meet ts1 ts2 |> function
@@ -480,7 +482,7 @@ module Type = struct
   let is_any tp_expr = equal tp_expr any
 
   let is_set tp_expr = match tp_expr with
-    | App (Set, _, _) -> true
+    | App (Map, [_; App(Bool, _, _)], _) -> true
     | _ -> false
   let is_ghost_var vdecl = vdecl.var_ghost
   let is_const_var vdecl = vdecl.var_const
@@ -498,7 +500,7 @@ module Type = struct
 
   
   let set_elem = function
-  | App (Set, elem :: _, _) -> elem
+  | App (Map, [elem; App (Bool, _, _)], _) -> elem
   | _ -> failwith "Expected Set type"
         
   let map_dom = function
@@ -762,11 +764,8 @@ module Expr = struct
     | ((Forall | Exists) as b), vs, trgs, e, _ ->
       Stdlib.Format.fprintf ppf "%s@ %a@ ::@ %a %a" (binder_to_string b)
       pr_var_decl_list vs pr_trgs trgs pr e
-    | Compr, vs, trgs, e, App (Set, _, _) ->
-        Stdlib.Format.fprintf ppf "{|@ @[%a@ ::@ %a@]@ |}" pr_var_decl_list vs
-          pr e
     | Compr, vs, trgs, e, _ ->
-        Stdlib.Format.fprintf ppf "[|@ @[%a@ ::@ %a@]@ |]" pr_var_decl_list vs
+        Stdlib.Format.fprintf ppf "{|@ @[%a@ ::@ %a@]@ |}" pr_var_decl_list vs
           pr e
 
   and pr_trgs ppf trgs = 
diff --git a/lib/backend/smtLibAST.ml b/lib/backend/smtLibAST.ml
index b89b34e..2f0eda0 100644
--- a/lib/backend/smtLibAST.ml
+++ b/lib/backend/smtLibAST.ml
@@ -112,7 +112,7 @@ let rec pr_sort ppf (sort : sort) =
   | App (Bool, [], _) -> fprintf ppf "Bool"
   | App (Ref, [], _) -> pr_smt_ident ppf PreambleConsts.loc_ident
   | App (Var qual_iden, [], _) -> pr_smt_ident ppf qual_iden
-  | App (Set, [ srt ], _) -> fprintf ppf "@[<2>(Set %a)@]" pr_sort srt
+  | App (Map, [ srt; App (Bool, _, _)], _) -> fprintf ppf "@[<2>(Set %a)@]" pr_sort srt
   | App (Map, [ srt1; srt2 ], _) ->
       fprintf ppf "@[<2>(Array %a %a)@]" pr_sort srt1 pr_sort srt2
   | App (Data (id, _), [], _) -> pr_smt_ident ppf id
diff --git a/lib/frontend/rewrites/atomicityAnalysis.ml b/lib/frontend/rewrites/atomicityAnalysis.ml
index 3166c16..d60e39f 100644
--- a/lib/frontend/rewrites/atomicityAnalysis.ml
+++ b/lib/frontend/rewrites/atomicityAnalysis.ml
@@ -118,15 +118,6 @@ let close_au ~loc token atomicity_state : atomicity_check =
       { atomicity_state with au_opened; atomic_step_taken = false }
     else { atomicity_state with au_opened }
 
-let rec is_expr_atomic (expr : Expr.t) : bool =
-  match expr with
-  | App (constr, expr_args, _) -> (
-      match constr with
-      | Null | Int _ | Real _ | Bool _ -> true
-      | Var _ -> ( match expr_args with [] -> true | _ -> false)
-      | Read | Cas -> is_expr_atomic (List.hd_exn expr_args)
-      | _ -> false)
-  | _ -> false
 
 let rewrite_au_cmnds (stmt : Stmt.t) : (Stmt.t, atomicity_check) Rewriter.t_ext
     =
@@ -256,16 +247,9 @@ let rewrite_au_cmnds (stmt : Stmt.t) : (Stmt.t, atomicity_check) Rewriter.t_ext
             let* _ = Rewriter.set_user_state atomicity_state in
             Rewriter.return stmt
     | Basic (Return return_expr) ->
-        if is_ghost_scope then
-          Error.error stmt.stmt_loc "Cannot return in a ghost block" (* TODO: move this to Typing *)
-        else if is_expr_atomic return_expr then
-          let atomicity_state = take_atomic_step ~loc atomicity_state in
-          let* _ = Rewriter.set_user_state atomicity_state in
-          Rewriter.return stmt
-        else
-          let atomicity_state = take_non_atomic_step ~loc atomicity_state in
-          let* _ = Rewriter.set_user_state atomicity_state in
-          Rewriter.return stmt
+        let atomicity_state = take_atomic_step ~loc atomicity_state in
+        let* _ = Rewriter.set_user_state atomicity_state in
+        Rewriter.return stmt
     | Basic (Use use_desc) -> (
         let* symbol = Rewriter.find_and_reify use_desc.use_name in
         match symbol with
diff --git a/lib/frontend/rewrites/rewrites.ml b/lib/frontend/rewrites/rewrites.ml
index 0e83446..90bbcc1 100644
--- a/lib/frontend/rewrites/rewrites.ml
+++ b/lib/frontend/rewrites/rewrites.ml
@@ -118,7 +118,7 @@ let rec rewrite_compr_expr (expr : expr) : expr Rewriter.t =
 
       let postcond =
         let spec_form =
-          if Type.is_set ret_typ then
+          (*if Type.is_set ret_typ then
             let var_decl = List.hd_exn v_l in
 
             Expr.mk_binder ~typ:Type.bool Forall [ var_decl ]
@@ -143,7 +143,7 @@ let rec rewrite_compr_expr (expr : expr) : expr Rewriter.t =
                        inner_expr;
                      ];
                  ])
-          else
+          else*)
             (* Type.is_map ret_typ *)
             let var_decl = List.hd_exn v_l in
             let lookup_expr =
diff --git a/lib/frontend/typing.ml b/lib/frontend/typing.ml
index 8662ddf..ac26ab8 100644
--- a/lib/frontend/typing.ml
+++ b/lib/frontend/typing.ml
@@ -55,7 +55,7 @@ module ProcessTypeExpr = struct
         | ModInst _ -> unexpected_functor_error tp_attr.type_loc
         | _ -> Error.type_error tp_attr.type_loc "Expected type identifier.")
     | App (Var _, _, tp_attr) -> unexpected_functor_error tp_attr.type_loc
-    | App (((Set | Fld) as constr), tp_list, tp_attr) -> (
+    | App ((Fld as constr), tp_list, tp_attr) -> (
         match tp_list with
         | [ tp_arg ] ->
             let+ tp_arg' = process_type_expr tp_arg in
@@ -705,7 +705,6 @@ let rec process_expr (expr : expr) (expected_typ : type_expr) : expr Rewriter.t
           
           let inner_expr_expected_typ =
             let ty = match expected_typ with
-            | App (Set, _, _) -> Type.bool
             | App (Map, [ _; tp ], _) -> tp
             | _ -> Type.any
             in ty |> Type.set_ghost_to expected_typ
@@ -1139,7 +1138,7 @@ module ProcessCallable = struct
         "Internal error: process_au_action_stmt called with non-callable \
          expression"
 
-  let rec process_basic_stmt (expected_return_type : Type.t)
+  let rec process_basic_stmt call_decl
       (basic_stmt : Stmt.basic_stmt_desc) (stmt_loc: Loc.t) (disam_tbl : DisambiguationTbl.t) :
       (Stmt.basic_stmt_desc * DisambiguationTbl.t) Rewriter.t =
     let open Rewriter.Syntax in
@@ -1192,9 +1191,8 @@ module ProcessCallable = struct
                 assign_is_init = true;
               }
           in
-          
           process_basic_stmt 
-            expected_return_type
+            call_decl
             (Assign assign_desc)
             stmt_loc
             disam_tbl'
@@ -1235,7 +1233,7 @@ module ProcessCallable = struct
                 field_read_is_init = assign_desc.assign_is_init;
               }
           in
-          process_basic_stmt expected_return_type (Stmt.FieldRead field_read_desc) stmt_loc disam_tbl
+          process_basic_stmt call_decl (Stmt.FieldRead field_read_desc) stmt_loc disam_tbl
         (* AU action *)
         | App (Var qual_ident, _, _) as assign_rhs when Predefs.is_qual_ident_au_cmnd qual_ident ->
           process_au_action_stmt assign_lhs var_decls_lhs assign_rhs stmt_loc disam_tbl
@@ -1453,6 +1451,10 @@ module ProcessCallable = struct
       let+ qual_ident, _ = get_assign_lhs qual_ident ~is_init:false in
       Stmt.Havoc qual_ident, disam_tbl
     | Return expr ->
+      let expected_return_type = Callable.return_type call_decl in
+      if is_ghost_scope && Poly.(call_decl.call_decl_kind = Proc) then
+        Error.type_error stmt_loc "Cannot return in a ghost block";
+
       let+ expr =
         disambiguate_process_expr expr expected_return_type disam_tbl
       in
@@ -1613,9 +1615,9 @@ module ProcessCallable = struct
       internal_error stmt_loc
         "Did not expect Fpu stmts in AST at this stage."
         
-  let process_stmt ?(new_scope = true) (expected_return_type : Type.t)
+  let process_stmt ?(new_scope = true) call_decl
       (stmt : Stmt.t) (disam_tbl : DisambiguationTbl.t) :
-      (Stmt.t * DisambiguationTbl.t) Rewriter.t =
+    (Stmt.t * DisambiguationTbl.t) Rewriter.t =
     let rec process_stmt ?(new_scope = true) stmt disam_tbl =
       Logs.debug (fun m -> m "process_stmt: %a" Stmt.pr stmt);
       let open Rewriter.Syntax in
@@ -1624,7 +1626,7 @@ module ProcessCallable = struct
         match stmt.Stmt.stmt_desc with
         | Basic basic_stmt ->
           let+ basic_stmt, disam_tbl' =
-            process_basic_stmt expected_return_type basic_stmt (Stmt.to_loc stmt) disam_tbl
+            process_basic_stmt call_decl basic_stmt (Stmt.to_loc stmt) disam_tbl
           in
           (Stmt.Basic basic_stmt, disam_tbl')
         | Block block_desc ->
@@ -1779,7 +1781,6 @@ module ProcessCallable = struct
 
           func_def
       | ProcDef proc_def ->
-          let expected_return_type = Callable.return_type call_decl in
           let+ proc_body =
             Rewriter.Option.map proc_def.proc_body ~f:(fun stmt ->
                 (* Logs.debug (fun m -> m "Typing.process_callable: Processing stmt: %a" Stmt.pr stmt); *)
@@ -1794,7 +1795,7 @@ module ProcessCallable = struct
                       (List.map disam_tbl ~f:Map.to_alist));
 
                 let+ stmt, _disam_tbl =
-                  process_stmt ~new_scope:false expected_return_type stmt
+                  process_stmt ~new_scope:false call_decl stmt
                     disam_tbl
                 in
                 stmt)
diff --git a/test/bugs/set_type.rav b/test/bugs/set_type.rav
deleted file mode 100644
index 34c21e5..0000000
--- a/test/bugs/set_type.rav
+++ /dev/null
@@ -1,3 +0,0 @@
-proc p() {
-  var x: Map[Ref, Bool] := {| l: Ref :: false |};
-}
\ No newline at end of file