diff --git a/Source/Concurrency/LinearRewriter.cs b/Source/Concurrency/LinearRewriter.cs
index 9666d807f..d7b41a311 100644
--- a/Source/Concurrency/LinearRewriter.cs
+++ b/Source/Concurrency/LinearRewriter.cs
@@ -155,7 +155,6 @@ public List<Cmd> RewriteCallCmd(CallCmd callCmd)
       case "Map_MakeEmpty":
       case "Map_Pack":
       case "Map_Unpack":
-      case "Map_Assume":
         return new List<Cmd>{callCmd};
       case "Set_Split":
         return RewriteSetSplit(callCmd);
@@ -177,6 +176,10 @@ public List<Cmd> RewriteCallCmd(CallCmd callCmd)
         return RewriteMapGet(callCmd);
       case "Map_Put":
         return RewriteMapPut(callCmd);
+      case "Map_GetValue":
+        return RewriteMapGetValue(callCmd);
+      case "Map_PutValue":
+        return RewriteMapPutValue(callCmd);
       default:
         Contract.Assume(false);
         return null;
@@ -481,6 +484,50 @@ private List<Cmd> RewriteMapPut(CallCmd callCmd)
     return cmdSeq;
   }
 
+  private List<Cmd> RewriteMapGetValue(CallCmd callCmd)
+  {
+    var cmdSeq = new List<Cmd>();
+    var path = callCmd.Ins[0];
+    var k = callCmd.Ins[1];
+    var v = callCmd.Outs[0];
+
+    var instantiation = monomorphizer.GetTypeInstantiation(callCmd.Proc);
+    var domain = instantiation["K"];
+    var range = instantiation["V"];
+    var mapContainsFunc = MapContains(domain, range);
+    var mapRemoveFunc = MapRemove(domain, range);
+    var mapAtFunc = MapAt(domain, range);
+    cmdSeq.Add(AssertCmd(callCmd.tok, ExprHelper.FunctionCall(mapContainsFunc, path, k), "Map_GetValue failed"));
+    var oneConstructor = OneConstructor(domain);
+    cmdSeq.Add(CmdHelper.AssignCmd(v.Decl, ExprHelper.FunctionCall(mapAtFunc, path, k)));
+    cmdSeq.Add(
+      CmdHelper.AssignCmd(CmdHelper.ExprToAssignLhs(path), ExprHelper.FunctionCall(mapRemoveFunc, path, k)));
+
+    ResolveAndTypecheck(options, cmdSeq);
+    return cmdSeq;
+  }
+
+  private List<Cmd> RewriteMapPutValue(CallCmd callCmd)
+  {
+    var cmdSeq = new List<Cmd>();
+    var path = callCmd.Ins[0];
+    var k = callCmd.Ins[1];
+    var v = callCmd.Ins[2];
+
+    var instantiation = monomorphizer.GetTypeInstantiation(callCmd.Proc);
+    var domain = instantiation["K"];
+    var range = instantiation["V"];
+    var mapContainsFunc = MapContains(domain, range);
+    var mapUpdateFunc = MapUpdate(domain, range);
+    var attribute = new QKeyValue(Token.NoToken, "linear", new List<object>(), null);
+    cmdSeq.Add(new AssumeCmd(Token.NoToken, Expr.Not(ExprHelper.FunctionCall(mapContainsFunc, path, k)), attribute));
+    cmdSeq.Add(
+      CmdHelper.AssignCmd(CmdHelper.ExprToAssignLhs(path), ExprHelper.FunctionCall(mapUpdateFunc, path, k, v)));
+
+    ResolveAndTypecheck(options, cmdSeq);
+    return cmdSeq;
+  }
+
   private void ResolveAndTypecheck(CoreOptions options, IEnumerable<Absy> absys)
   {
     var rc = new ResolutionContext(null, options);
diff --git a/Source/Concurrency/LinearTypeChecker.cs b/Source/Concurrency/LinearTypeChecker.cs
index e17df72a7..36a6e93f2 100644
--- a/Source/Concurrency/LinearTypeChecker.cs
+++ b/Source/Concurrency/LinearTypeChecker.cs
@@ -462,6 +462,10 @@ public override Cmd VisitAssignCmd(AssignCmd node)
           {
             Error(rhs, $"linear variable {rhs.Decl.Name} can occur at most once as the source of an assignment");
           }
+          else if (InvalidAssignmentWithKeyCollection(lhs.DeepAssignedVariable, rhs.Decl))
+          {
+            Error(rhs, $"Mismatch in key collection between source and target");
+          }
           else
           {
             rhsVars.Add(rhs.Decl);
@@ -486,6 +490,10 @@ public override Cmd VisitAssignCmd(AssignCmd node)
             {
               Error(arg, $"linear variable {ie.Decl.Name} can occur at most once as the source of an assignment");
             }
+            else if (InvalidAssignmentWithKeyCollection(field, ie.Decl))
+            {
+              Error(arg, $"Mismatch in key collection between source and target");
+            }
             else
             {
               rhsVars.Add(ie.Decl);
@@ -568,6 +576,11 @@ public override Cmd VisitCallCmd(CallCmd node)
           Error(node, $"linear variable {actual.Decl.Name} can occur only once as an input parameter");
           continue;
         }
+        if (!isPrimitive && InvalidAssignmentWithKeyCollection(formal, actual.Decl))
+        {
+          Error(node, $"Mismatch in key collection between source and target");
+          continue;
+        }
         inVars.Add(actual.Decl);
         if (actual.Decl is GlobalVariable && actualKind == LinearKind.LINEAR_IN)
         {
@@ -590,6 +603,11 @@ public override Cmd VisitCallCmd(CallCmd node)
           Error(node, $"only linear parameter can be assigned to a linear variable: {formal}");
           continue;
         }
+        if (!isPrimitive && InvalidAssignmentWithKeyCollection(actual.Decl, formal))
+        {
+          Error(node, $"Mismatch in key collection between source and target");
+          continue;
+        }
       }
 
       var globalOutVars = node.Outs.Select(ie => ie.Decl).ToHashSet();
@@ -598,6 +616,8 @@ public override Cmd VisitCallCmd(CallCmd node)
         Error(node, $"global variable passed as input to pure call but not received as output: {v}");
       });
 
+      var originalProc = (Procedure)Monomorphizer.GetOriginalDecl(node.Proc);
+
       if (isPrimitive)
       {
         var modifiedArgument = CivlPrimitives.ModifiedArgument(node)?.Decl;
@@ -612,17 +632,48 @@ public override Cmd VisitCallCmd(CallCmd node)
             Error(node, $"primitive assigns to input variable that is also an output variable: {modifiedArgument}");
           }
           else if (modifiedArgument is GlobalVariable &&
-                   enclosingProc is not YieldProcedureDecl &&
-                   enclosingProc.Modifies.All(v => v.Decl != modifiedArgument))
+                    enclosingProc is not YieldProcedureDecl &&
+                    enclosingProc.Modifies.All(v => v.Decl != modifiedArgument))
           {
             var str = enclosingProc is ActionDecl ? "action's" : "procedure's";
             Error(node,
               $"primitive assigns to a global variable that is not in the enclosing {str} modifies clause: {modifiedArgument}");
           }
+
+          if (originalProc.Name == "Map_Split")
+          {
+            if (InvalidAssignmentWithKeyCollection(node.Outs[0].Decl, modifiedArgument))
+            {
+              Error(node, $"Mismatch in key collection between source and target");
+            }
+          }
+          else if (originalProc.Name == "Map_Join")
+          {
+            if (node.Ins[1] is IdentifierExpr ie)
+            {
+              if (InvalidAssignmentWithKeyCollection(modifiedArgument, ie.Decl))
+              {
+                Error(node, $"Mismatch in key collection between source and target");
+              }
+            }
+          }
+          else if (originalProc.Name == "Map_Get" || originalProc.Name == "Map_Put")
+          {
+            if (!AreKeysCollected(modifiedArgument))
+            {
+              Error(node, $"Keys must be collected");
+            }
+          }
+          else if (originalProc.Name == "Map_GetValue" || originalProc.Name == "Map_PutValue")
+          {
+            if (AreKeysCollected(modifiedArgument))
+            {
+              Error(node, $"Keys must not be collected");
+            }
+          }
         }
       }
 
-      var originalProc = (Procedure)Monomorphizer.GetOriginalDecl(node.Proc);
       if (originalProc.Name == "create_multi_asyncs" || originalProc.Name == "create_asyncs")
       {
         var actionDecl = GetActionDeclFromCreateAsyncs(node);
@@ -749,6 +800,25 @@ public override Variable VisitVariable(Variable node)
       return node;
     }
 
+    private bool AreKeysCollected(Variable v)
+    {
+      var attr = QKeyValue.FindAttribute(v.Attributes, x => x.Key == "linear");
+      var attrParams = attr == null ? new List<object>() : attr.Params;
+      foreach (var param in attrParams)
+      {
+        if (param is string s && s == "no_collect_keys")
+        {
+          return false;
+        }
+      }
+      return true;
+    }
+
+    private bool InvalidAssignmentWithKeyCollection(Variable target, Variable source)
+    {
+      return AreKeysCollected(target) && !AreKeysCollected(source);
+    }
+
     private void CheckLinearStoreAccessInGuards()
     {
       program.Implementations.ForEach(impl => {
diff --git a/Source/Core/CivlAttributes.cs b/Source/Core/CivlAttributes.cs
index 917d00c2f..42b25efab 100644
--- a/Source/Core/CivlAttributes.cs
+++ b/Source/Core/CivlAttributes.cs
@@ -218,7 +218,8 @@ public static class CivlPrimitives
     public static HashSet<string> LinearPrimitives = new()
     {
       "Loc_New", "KeyedLocSet_New",
-      "Map_MakeEmpty", "Map_Pack", "Map_Unpack", "Map_Split", "Map_Join", "Map_Get", "Map_Put", "Map_Assume",
+      "Map_MakeEmpty", "Map_Pack", "Map_Unpack", "Map_Split", "Map_Join",
+      "Map_Get", "Map_Put", "Map_GetValue", "Map_PutValue",
       "Set_MakeEmpty", "Set_Split", "Set_Get", "Set_Put", "One_Split", "One_Get", "One_Put"
     };
 
@@ -248,7 +249,6 @@ public static IdentifierExpr ModifiedArgument(CallCmd callCmd)
         case "Map_MakeEmpty":
         case "Map_Pack":
         case "Map_Unpack":
-        case "Map_Assume":
           return null;
         default:
           return ExtractRootFromAccessPathExpr(callCmd.Ins[0]);
diff --git a/Source/Core/base.bpl b/Source/Core/base.bpl
index cd3d482ea..5185f5a43 100644
--- a/Source/Core/base.bpl
+++ b/Source/Core/base.bpl
@@ -319,17 +319,15 @@ pure procedure {:inline 1} Map_Pack<K,V>({:linear_in} dom: Set K, val: [K]V) ret
 }
 pure procedure {:inline 1} Map_Unpack<K,V>({:linear_in} m: Map K V) returns ({:linear} dom: Set K, val: [K]V)
 {
- dom := m->dom;
- val := m->val;
+  dom := m->dom;
+  val := m->val;
 }
 pure procedure Map_Split<K,V>({:linear} path: Map K V, s: Set K) returns ({:linear} m: Map K V);
 pure procedure Map_Join<K,V>({:linear} path: Map K V, {:linear_in} m: Map K V);
 pure procedure Map_Get<K,V>({:linear} path: Map K V, k: K) returns ({:linear} l: One K, {:linear} v: V);
 pure procedure Map_Put<K,V>({:linear} path: Map K V, {:linear_in} l: One K, {:linear_in} v: V);
-pure procedure {:inline 1} Map_Assume<K,V>({:linear} src: Map K V, {:linear} dst: Map K V)
-{
-  assume Set_IsDisjoint(src->dom, dst->dom);
-}
+pure procedure Map_GetValue<K,V>({:linear} path: Map K V, k: K) returns ({:linear} v: V);
+pure procedure Map_PutValue<K,V>({:linear} path: Map K V, k: K, {:linear_in} v: V);
 
 type Loc;
 
diff --git a/Test/civl/samples/treiber-stack.bpl b/Test/civl/samples/treiber-stack.bpl
index 4db45658e..b9a03d881 100644
--- a/Test/civl/samples/treiber-stack.bpl
+++ b/Test/civl/samples/treiber-stack.bpl
@@ -14,7 +14,7 @@ datatype Treiber<T> { Treiber(top: Option LocTreiberNode, {:linear} stack: Stack
 type X; // module type parameter
 
 var {:layer 4, 5} Stack: Map Loc (Vec X);
-var {:layer 0, 4} {:linear} ts: Map Loc (Treiber X);
+var {:layer 0, 4} {:linear "no_collect_keys"} ts: Map Loc (Treiber X);
 
 /// Yield invariants
 
@@ -61,31 +61,31 @@ invariant (var t := ts->val[loc_t]; Map_At(t->stack, new_loc_n) == node && !Betw
 
 /// Layered implementation
 
-atomic action {:layer 5} AtomicAlloc() returns (loc_t: Loc)
+atomic action {:layer 5} AtomicAlloc() returns ({:linear} one_loc_t: One Loc)
 modifies Stack;
 {
-  var {:linear} one_loc_t: One Loc;
+  var loc_t: Loc;
 
   call one_loc_t := Loc_New();
   loc_t := one_loc_t->val;
   assume !Map_Contains(Stack, loc_t);
   Stack := Map_Update(Stack, loc_t, Vec_Empty());
 }
-yield procedure {:layer 4} Alloc() returns (loc_t: Loc)
+yield procedure {:layer 4} Alloc() returns ({:linear} one_loc_t: One Loc)
 refines AtomicAlloc;
 preserves call StackDom();
 {
   var top: Option LocTreiberNode;
   var {:linear} stack: StackMap X;
   var {:linear} treiber: Treiber X;
-  var {:linear} one_loc_t: One Loc;
+  var loc_t: Loc;
 
   top := None();
   call stack := Map_MakeEmpty();
   treiber := Treiber(top, stack);
   call one_loc_t := Loc_New();
   loc_t := one_loc_t->val;
-  call AllocTreiber#0(one_loc_t, treiber);
+  call AllocTreiber#0(loc_t, treiber);
   call {:layer 4} Stack := Copy(Map_Update(Stack, loc_t, Vec_Empty()));
   call {:layer 4} AbsLemma(treiber);
 }
@@ -161,7 +161,6 @@ atomic action {:layer 4} AtomicAllocNode#3(loc_t: Loc, x: X)
 modifies ts;
 asserts Map_Contains(ts, loc_t);
 {
-  var {:linear} one_loc_t: One Loc;
   var {:linear} treiber: Treiber X;
   var top: Option LocTreiberNode;
   var {:linear} stack: StackMap X;
@@ -169,7 +168,7 @@ asserts Map_Contains(ts, loc_t);
   var {:linear} loc_pieces: Set LocTreiberNode;
   var {:linear} left_loc_piece: One LocTreiberNode;
   
-  call one_loc_t, treiber := Map_Get(ts, loc_t);
+  call treiber := Map_GetValue(ts, loc_t);
   Treiber(top, stack) := treiber;
   assume loc_n is None || Map_Contains(stack, loc_n->t);
   call loc, loc_pieces := KeyedLocSet_New(AllLocPieces());
@@ -178,7 +177,7 @@ asserts Map_Contains(ts, loc_t);
   call right_loc_piece := One_Get(loc_pieces, KeyedLoc(loc, Right()));
   call Map_Put(stack, left_loc_piece, Node(loc_n, x));
   treiber := Treiber(top, stack);
-  call Map_Put(ts, one_loc_t, treiber);
+  call Map_PutValue(ts, loc_t, treiber);
 }
 yield procedure {:layer 3} AllocNode#3(loc_t: Loc, x: X)
   returns (loc_n: Option LocTreiberNode, new_loc_n: LocTreiberNode, {:linear} right_loc_piece: One LocTreiberNode)
@@ -201,7 +200,6 @@ refines AtomicAllocNode#3;
 atomic action {:layer 4} AtomicPopIntermediate(loc_t: Loc) returns (success: bool, x_opt: Option X)
 modifies ts;
 {
-  var {:linear} one_loc_t: One Loc;
   var loc_n: Option LocTreiberNode;
   var {:linear} treiber: Treiber X;
   var top: Option LocTreiberNode;
@@ -210,12 +208,12 @@ modifies ts;
 
   assert Map_Contains(ts, loc_t);
   if (*) {
-    call one_loc_t, treiber := Map_Get(ts, loc_t);
+    call treiber := Map_GetValue(ts, loc_t);
     Treiber(top, stack) := treiber;
     assume top is Some && Map_Contains(stack, top->t);
     Node(loc_n, x) := Map_At(stack, top->t);
     treiber := Treiber(loc_n, stack);
-    call Map_Put(ts, one_loc_t, treiber);
+    call Map_PutValue(ts, loc_t, treiber);
     x_opt := Some(x);
     success := true;
   }
@@ -293,18 +291,17 @@ atomic action {:layer 1} AtomicLoadNode#0(loc_t: Loc, loc_n: LocTreiberNode) ret
 modifies ts;
 refines AtomicLoadNode#1;
 {
-  var {:linear} one_loc_t: One Loc;
   var {:linear} treiber: Treiber X;
   var top: Option LocTreiberNode;
   var {:linear} stack: StackMap X;
   var {:linear} one_loc_n: One LocTreiberNode;
 
-  call one_loc_t, treiber := Map_Get(ts, loc_t);
+  call treiber := Map_GetValue(ts, loc_t);
   Treiber(top, stack) := treiber;
   call one_loc_n, node := Map_Get(stack, loc_n);
   call Map_Put(stack, one_loc_n, node);
   treiber := Treiber(top, stack);
-  call Map_Put(ts, one_loc_t, treiber);
+  call Map_PutValue(ts, loc_t, treiber);
 }
 yield procedure {:layer 0} LoadNode#0(loc_t: Loc, loc_n: LocTreiberNode) returns (node: StackElem X);
 refines AtomicLoadNode#0;
@@ -312,12 +309,11 @@ refines AtomicLoadNode#0;
 atomic action {:layer 1,2} AtomicReadTopOfStack#0(loc_t: Loc) returns (loc_n: Option LocTreiberNode)
 modifies ts;
 {
-  var {:linear} one_loc_t: One Loc;
   var {:linear} treiber: Treiber X;
 
-  call one_loc_t, treiber := Map_Get(ts, loc_t);
+  call treiber := Map_GetValue(ts, loc_t);
   loc_n := treiber->top;
-  call Map_Put(ts, one_loc_t, treiber);
+  call Map_PutValue(ts, loc_t, treiber);
 }
 yield procedure {:layer 0} ReadTopOfStack#0(loc_t: Loc) returns (loc_n: Option LocTreiberNode);
 refines AtomicReadTopOfStack#0;
@@ -327,16 +323,15 @@ atomic action {:layer 1,4} AtomicWriteTopOfStack#0(
 modifies ts;
 {
   var {:linear} treiber: Treiber X;
-  var {:linear} one_loc_t: One Loc;
 
-  call one_loc_t, treiber := Map_Get(ts, loc_t);
+  call treiber := Map_GetValue(ts, loc_t);
   if (old_loc_n == treiber->top) {
     treiber->top := new_loc_n;
     success := true;
   } else {
     success := false;
   }
-  call Map_Put(ts, one_loc_t, treiber);
+  call Map_PutValue(ts, loc_t, treiber);
 }
 yield procedure {:layer 0} WriteTopOfStack#0(
   loc_t: Loc, old_loc_n: Option LocTreiberNode, new_loc_n: Option LocTreiberNode) returns (success: bool);
@@ -345,22 +340,21 @@ refines AtomicWriteTopOfStack#0;
 atomic action {:layer 1,3} AtomicAllocNode#0(loc_t: Loc, {:linear_in} loc_piece: One LocTreiberNode, node: StackElem X)
 modifies ts;
 {
-  var {:linear} one_loc_t: One Loc;
   var {:linear} treiber: Treiber X;
 
-  call one_loc_t, treiber := Map_Get(ts, loc_t);
+  call treiber := Map_GetValue(ts, loc_t);
   call Map_Put(treiber->stack, loc_piece, node);
-  call Map_Put(ts, one_loc_t, treiber);
+  call Map_PutValue(ts, loc_t, treiber);
 }
 yield procedure {:layer 0} AllocNode#0(loc_t: Loc, {:linear_in} loc_piece: One LocTreiberNode, node: StackElem X);
 refines AtomicAllocNode#0;
 
-atomic action {:layer 1,4} AtomicAllocTreiber#0({:linear_in} one_loc_t: One Loc, {:linear_in} treiber: Treiber X)
+atomic action {:layer 1,4} AtomicAllocTreiber#0(loc_t: Loc, {:linear_in} treiber: Treiber X)
 modifies ts;
 {
-  call Map_Put(ts, one_loc_t, treiber);
+  call Map_PutValue(ts, loc_t, treiber);
 }
-yield procedure {:layer 0} AllocTreiber#0({:linear_in} one_loc_t: One Loc, {:linear_in} treiber: Treiber X);
+yield procedure {:layer 0} AllocTreiber#0(loc_t: Loc, {:linear_in} treiber: Treiber X);
 refines AtomicAllocTreiber#0;
 
 /// Proof of abstraction with a manual encoding of termination