Update code according to @mwiencek comments.

Please do not merge, i'll rebase after we finish with this.