This should be the last step in the current cleanup. Follow-ups should resolve the TODO about cost calc and enable the more general case where we extract different elements.