Perhaps you mean antanaclasis:
The repetition of a word or phrase whose meaning changes in the second instance.
From Silva Rhetoricae, "antanaclasis".
Some examples from that site:
Your argument is sound...all sound. —Benjamin Franklin
If you aren't fired with enthusiasm, you will be fired with enthusiasm. —Vince Lombardi
The overarching figure for such repetition in brief sentences is paregmenon:
A general term for the repetition of a word or its cognates in a short sentence.
Polyptoton (as in your first example), is also a possibility:
Repeating a word, but in a different form. Using a cognate of a given word in close proximity.
Adnominatio is sometimes used as a synonym of polyptoton.
The repetition of a word in a contrary sense.
is sometimes used as a synonym of antanaclasis.
(All definitions from Silva Rhetoricae.)