Use statement walker for constant folding.

- makes the code compact (gets rid of MLFunction walking logic)
- makes it natural to extend to fold affine map loop bounds
  and if conditions (upcoming CL)

PiperOrigin-RevId: 214668957