If that is really the only function that causes problems by blocking then you could simply take the pragmatic approach and call a function from inside the rainbow function that reads the button input and return from the rainbow function if the button count changes
NOTE : I do not mean that you should use an interrupt