Fix corner cases for stack shrink