Move variable assignment handling to generation context

This allows the parsing code to be cleaner, as it
doesn't have to care about variable assignments.

Bug: 228518745
Test: go test
Change-Id: I33425c2fb51acab4901bfa82a53d337b75210f8e