Rewrite context generation to be more concise and correct

This handles cases where a higher kind is not the first argument (e.g. in * -> (* -> *) -> *), which would previously not get generated correctly.