diff --git a/src/analyses/memLeak.ml b/src/analyses/memLeak.ml
index dbaa2d69fc..9c09c05cf6 100644
--- a/src/analyses/memLeak.ml
+++ b/src/analyses/memLeak.ml
@@ -14,9 +14,10 @@ struct
   let name () = "memLeak"
 
   module D = ToppedVarInfoSet
-  module C = Lattice.Unit
+  module C = D
+  module P = IdentityP (D)
 
-  let context _ _ = ()
+  let context _ d = d
 
   (* HELPER FUNCTIONS *)
   let warn_for_multi_threaded ctx =
diff --git a/src/witness/myARG.ml b/src/witness/myARG.ml
index 068aed7a22..373a66d3d6 100644
--- a/src/witness/myARG.ml
+++ b/src/witness/myARG.ml
@@ -141,7 +141,7 @@ struct
   let equal_node_context _ _ = failwith "StackNode: equal_node_context"
 end
 
-module Stack (Cfg:CfgForward) (Arg: S):
+module Stack (Arg: S with module Edge = InlineEdge):
   S with module Node = StackNode (Arg.Node) and module Edge = Arg.Edge =
 struct
   module Node = StackNode (Arg.Node)
@@ -156,45 +156,30 @@ struct
     | n :: stack ->
       let cfgnode = Arg.Node.cfgnode n in
       match cfgnode with
-      | Function _ -> (* TODO: can this be done without Cfg? *)
+      | Function _ -> (* TODO: can this be done without cfgnode? *)
         begin match stack with
           (* | [] -> failwith "StackArg.next: return stack empty" *)
           | [] -> [] (* main return *)
           | call_n :: call_stack ->
-            let call_cfgnode = Arg.Node.cfgnode call_n in
             let call_next =
-              Cfg.next call_cfgnode
+              Arg.next call_n
               (* filter because infinite loops starting with function call
                  will have another Neg(1) edge from the head *)
-              |> List.filter (fun (locedges, to_node) ->
-                  List.exists (function
-                      | (_, Proc _) -> true
-                      | (_, _) -> false
-                    ) locedges
+              |> List.filter_map (fun (edge, to_n) ->
+                  match edge with
+                  | InlinedEdge _ -> Some to_n
+                  | _ -> None
                 )
             in
-            begin match call_next with
-              | [] -> failwith "StackArg.next: call next empty"
-              | [(_, return_node)] ->
-                begin match Arg.Node.move_opt call_n return_node with
-                  (* TODO: Is it possible to have a calling node without a returning node? *)
-                  (* | None -> [] *)
-                  | None -> failwith "StackArg.next: no return node"
-                  | Some return_n ->
-                    (* TODO: Instead of next & filter, construct unique return_n directly. Currently edge missing. *)
-                    Arg.next n
-                    |> List.filter (fun (edge, to_n) ->
-                        (* let to_cfgnode = Arg.Node.cfgnode to_n in
-                        MyCFG.Node.equal to_cfgnode return_node *)
-                        Arg.Node.equal_node_context to_n return_n
-                      )
-                    |> List.map (fun (edge, to_n) ->
-                        let to_n' = to_n :: call_stack in
-                        (edge, to_n')
-                      )
-                end
-              | _ :: _ :: _ -> failwith "StackArg.next: call next ambiguous"
-            end
+            Arg.next n
+            |> List.filter_map (fun (edge, to_n) ->
+                if BatList.mem_cmp Arg.Node.compare to_n call_next then (
+                  let to_n' = to_n :: call_stack in
+                  Some (edge, to_n')
+                )
+                else
+                  None
+              )
         end
       | _ ->
         let+ (edge, to_n) = Arg.next n in
diff --git a/src/witness/svcomp.ml b/src/witness/svcomp.ml
index eae97b1199..89487ea8d4 100644
--- a/src/witness/svcomp.ml
+++ b/src/witness/svcomp.ml
@@ -76,9 +76,9 @@ sig
   val is_sink: Arg.Node.t -> bool
 end
 
-module StackTaskResult (Cfg:MyCFG.CfgForward) (TaskResult: TaskResult) =
+module StackTaskResult (TaskResult: TaskResult with module Arg.Edge = MyARG.InlineEdge) =
 struct
-  module Arg = MyARG.Stack (Cfg) (TaskResult.Arg)
+  module Arg = MyARG.Stack (TaskResult.Arg)
 
   let result = TaskResult.result
 
diff --git a/src/witness/witness.ml b/src/witness/witness.ml
index d94960d542..235461c348 100644
--- a/src/witness/witness.ml
+++ b/src/witness/witness.ml
@@ -14,7 +14,7 @@ let write_file filename (module Task:Task) (module TaskResult:WitnessTaskResult)
 
   let module TaskResult =
     (val if get_bool "witness.graphml.stack" then
-        (module StackTaskResult (Task.Cfg) (TaskResult) : WitnessTaskResult)
+        (module StackTaskResult (TaskResult) : WitnessTaskResult)
       else
         (module TaskResult)
     )
diff --git a/tests/regression/56-witness/52-witness-lifter-ps2.c b/tests/regression/56-witness/52-witness-lifter-ps2.c
new file mode 100644
index 0000000000..bcb7c1410c
--- /dev/null
+++ b/tests/regression/56-witness/52-witness-lifter-ps2.c
@@ -0,0 +1,35 @@
+// PARAM: --enable ana.sv-comp.enabled --enable ana.sv-comp.functions --enable witness.graphml.enabled --set ana.specification 'CHECK( init(main()), LTL(G valid-memtrack) )' --set ana.activated[+] memLeak --set ana.path_sens[+] memLeak --set ana.malloc.unique_address_count 1
+struct _twoIntsStruct {
+   int intOne ;
+   int intTwo ;
+};
+
+typedef struct _twoIntsStruct twoIntsStruct;
+
+void printStructLine(twoIntsStruct const *structTwoIntsStruct)
+{
+  return;
+}
+
+
+int main(int argc, char **argv)
+{
+  twoIntsStruct *data;
+  int tmp_1;
+
+
+  if (tmp_1 != 0) {
+    twoIntsStruct *dataBuffer = malloc(800UL);
+    data = dataBuffer;
+  }
+  else {
+
+    twoIntsStruct *dataBuffer_0 = malloc(800UL);
+    data = dataBuffer_0;
+  }
+
+  printStructLine((twoIntsStruct const *)data);
+  free((void *)data);
+
+  return;
+}
diff --git a/tests/regression/56-witness/53-witness-lifter-ps3.c b/tests/regression/56-witness/53-witness-lifter-ps3.c
new file mode 100644
index 0000000000..06b73b3888
--- /dev/null
+++ b/tests/regression/56-witness/53-witness-lifter-ps3.c
@@ -0,0 +1,39 @@
+// PARAM: --enable ana.sv-comp.enabled --enable ana.sv-comp.functions --enable witness.graphml.enabled --set ana.specification 'CHECK( init(main()), LTL(G valid-memtrack) )' --set ana.activated[+] memLeak --set ana.path_sens[+] memLeak --set ana.malloc.unique_address_count 1
+struct _twoIntsStruct {
+   int intOne ;
+   int intTwo ;
+};
+
+typedef struct _twoIntsStruct twoIntsStruct;
+
+void printStructLine(twoIntsStruct const *structTwoIntsStruct)
+{
+  return;
+}
+
+twoIntsStruct *foo() {
+  twoIntsStruct *data;
+  int tmp_1;
+
+  if (tmp_1 != 0) {
+    twoIntsStruct *dataBuffer = malloc(800UL);
+    data = dataBuffer;
+  }
+  else {
+
+    twoIntsStruct *dataBuffer_0 = malloc(800UL);
+    data = dataBuffer_0;
+  }
+  return data;
+}
+
+int main(int argc, char **argv)
+{
+  twoIntsStruct *data;
+  data = foo();
+
+  printStructLine((twoIntsStruct const *)data);
+  free((void *)data);
+
+  return;
+}