1 files changed, 48 insertions, 3 deletions
diff --git a/src/compiler/analyzer.zc b/src/compiler/analyzer.zc
index ff6a6c1..55a96bb 100644
--- a/src/compiler/analyzer.zc
+++ b/src/compiler/analyzer.zc
@@ -122,7 +122,7 @@ int shadow_finish(int from) {
   while (k < 24) { if ((e[2] >> k) & 1) flags = flags | effect_get(e[1], k); k++; }
   if ((flags & (EF_ALIAS | EF_IDENTITY)) && !(flags & EF_UNKNOWN) && !dry) {
    tunpack(e[3]); curinc = e[4]; panic = 0;
-   errat(); ended_message(e[6], e[7], e[8], e[9]); rt.put_s("; "); rt.put_s(sname[e[1]]); rt.put_s(", defined after this call, returns a borrow of its argument"); errend();
+   errat(); ended_message(e[6], e[7], e[8], e[9], 0); rt.put_s("; "); rt.put_s(sname[e[1]]); rt.put_s(", defined after this call, returns a borrow of its argument"); errend();
   }
   e = next;
  }
@@ -1009,7 +1009,7 @@ int end_owner(int i) {
  end_why = 0; end_whom = 0;
  local_ended[i] = info[2]; local_end_info[i] = info; lstate[i] = S_DEAD;
  while (j < nloc) {
-  if (j != i && isalias(j) && lroot(j) == i) { local_ended[j] = info[2]; local_end_info[j] = info; lstate[j] = S_DEAD; }
+  if (j != i && lstate[j] == S_CHILD && lroot(j) == i) { local_ended[j] = info[2]; local_end_info[j] = info; lstate[j] = S_DEAD; }	// aliases and owned children go with it
   j++;
  }
  return 0;
@@ -1147,6 +1147,7 @@ int block_close(int save) {
 	unused_check(save);
 	exit_vars(fdepth, 1);
 	lend_drop(fdepth, save);
+	block_end_owners(save);
 	alias_end(fdepth, save);
 	while (i < save) { if (lhome[i] >= fdepth) { lstate[i] = S_DEAD; lhome[i] = fdepth - 1; } i++; }
 	local_drop(save);
@@ -1155,6 +1156,26 @@ int block_close(int save) {
 	return 0;
 }
 
+// owners freed at this block's end: an outer alias of one, or an outer owner
+// stored into one, dies with it and remembers where (why 9).
+int block_end_owners(int save) {
+	int i = save, j; word info = 0;
+	if (!fchk || !live) return 0;
+	while (i < nloc) {
+		if (lhome[i] >= fdepth && lstate[i] == S_M && lown[i]) {
+			info = 0; j = lbase;
+			while (j < save) {
+				if (lstate[j] == S_CHILD && lroot(j) == i) {
+					if (!info) { info = talloc(6); info[0] = lname[i]; info[1] = lpos[i]; info[2] = tpack(); info[3] = 9; info[4] = 0; info[5] = info[2]; }
+					local_ended[j] = info[2]; local_end_info[j] = info; lstate[j] = S_DEAD;
+				}
+				j++;
+			}
+		}
+		i++;
+	}
+	return 0;
+}
 // an alias in an outer block whose owner is about to disappear dies with it
 int alias_end(int dep, int save) {
 	int i, r;
@@ -1532,6 +1553,19 @@ int assigncheck(int k, int i, int rc) {
 			else lstate[i] = S_PLAIN;
 			return C_PLAIN;
 		}
+		// an outer owner moved into a deeper container takes the container back:
+		// head = item after item.next = head. the container leaves its block.
+		if (st == S_CHILD && lown[i] && j >= 0 && lroot(i) == j && lstate[j] == S_M && lown[j]
+			&& lhome[j] >= lhome[i] && !local_captured[i] && !local_captured[j]) {
+			int k = lbase;
+			while (k < nloc) { if (k != i && lstate[k] == S_CHILD && lowner[k] == j) lowner[k] = i; k++; }
+			lstate[i] = S_M; lused[i] = 0;
+			lnsz[i] = lnsz[j]; lleaf[i] = lleaf[j];
+			local_ended[i] = 0; local_end_info[i] = 0;
+			end_why = 3; end_whom = lname[i]; end_owner(j); luse(j); ev_put(j, EV_MOVE, i);
+			if (ltag[i] & 15) lstate[i] = ltag[i] & 15;
+			return C_PLAIN;
+		}
 		if (!empty(st)) {
 			cerrn(i, "still owns an object; overwriting it would leak it");
 			return C_PLAIN;
@@ -1578,6 +1612,9 @@ int assigncheck(int k, int i, int rc) {
 		touch(i);
 		if (rc == C_FRESH) { emitown(i); return C_PLAIN; }
 		if (j < 0) return C_PLAIN;
+		// a back link: storing the container's own owner (or the container itself)
+		// into it. it stays a plain reference; ownership would form a cycle.
+		if (r == j) { touch(j); return C_PLAIN; }
 		js = lstate[j];
 		if (js == S_CHILD) cerrn(j, "already belongs to a container");
 		else if (js == S_BORROW) cerrn(j, "is borrowed and can't be stored");
@@ -1587,7 +1624,7 @@ int assigncheck(int k, int i, int rc) {
 			}
 			else { end_why = 4; end_whom = lname[i]; end_owner(j); luse(j); ev_put(j, EV_STORE, i); emitown(i); }	// storing it gives it away
 		} else if (js != st) cerrn(j, "its tag doesn't match its container's (both have to be m, u or d)");
-		else if (lhome[j] < lhome[r]) cerrn(j, "lives in an outer block than the container it's stored into");
+		else if (lhome[j] < lhome[r] && (r != i || !lown[j] || (local_captured && (local_captured[i] || local_captured[j])))) cerrn(j, "lives in an outer block than the container it's stored into");
 		else { lstate[j] = S_CHILD; lowner[j] = i; luse(j); ev_put(j, EV_STORE, i); emitown(i); }
 		return C_PLAIN;
 	}
@@ -1607,6 +1644,7 @@ int assigncheck(int k, int i, int rc) {
 // the function has no .u (its callers compiled after this point see it). an m
 // object is lent, not given, so returning it is an error, and so is a plain word
 // from an owned function that was called before its body.
+int ret_borrowed;
 int retcheck(int c) {
  int i;
  if (isptr(xty) || xty == tword) {
@@ -1633,9 +1671,16 @@ int retcheck(int c) {
 	if ((isptr(xty) || isagg(xty)) && ((xaf && xaf[2] == 4) || (xbor >= 0 && lafact[xbor] && lafact[xbor][2] == 4))) { cerr("an array literal cannot escape its function through return"); return 0; }
 	i = -1;
 	if (c >= C_VAR && tracked(c - C_VAR)) i = c - C_VAR;
+	// a value that comes from a parameter the caller still owns (not a .d one,
+	// which is tracked) is a borrow: it can't become the caller's owned result.
+	if (c != C_FRESH && i < 0 && isptr(xty) && effect_mask(c)) {
+		if (fret) { cerr("a borrowed parameter cannot be returned as an owned result"); return 0; }
+		ret_borrowed = 1;
+	}
 	if (i >= 0 && lstate[lroot(i)] == S_DEAD) { cerrn(i, "used after it was given away"); return 0; }
 	if (i >= 0 && (lstate[i] == S_D || (lstate[i] == S_M && lown[i])) && (fret || finfer)) lstate[i] = S_U;	// returning an owned object (a d object, or a local that owns one) gives it up, as u*x does
 	if (!fret && finfer && (c == C_FRESH || (i >= 0 && lstate[i] == S_U))) {
+		if (ret_borrowed) cerr("this function also returns a borrowed parameter, so its result can't be owned");
 		if (sflag[finsym] & F_EARLY) serrn(sname[finsym], "is called before its body, so its owned return can't be known there; declare the return type with .u");
 		ssig[finsym] = ssig[finsym] | 1;	// the return is owned: callers read it from the signature
 		sflag[finsym] = sflag[finsym] | F_RETOWN;