Lazy state: improve clarity

[josephjclark]

We need to improve clarity on a few things in lazy state

• Lazy state only works at the top level (not inside callbacks). I don't think we'll enforce this (I don't know how well we can?!) but we can document it at least. We should also make sure that the compiler stops traversal inside a new scope.
• Set up some clear wording that lazy state basically only applies to the parent function - it won't smartly walk up to the nearest operation. See discussion at Lazy state: doesn't work in non-operation function arguments #745 (comment)

This means updating docs, and maybe some unit tests in this repo