Sachverhalt:
A. Am 22. Juli 2005 wurde X._ festgenommen, nachdem A._ ihn beschuldigte hatte, sie am 17. Juli vergewaltigt zu haben. Am folgenden Tag wurde A._ mittels Videobefragung einvernommen. Im Anschluss daran wurde X._ zum Vergewaltigungsvorwurf befragt und freigelassen.
B. Vom 26. August bis 6. September 2005 war X._ in Untersuchungshaft wegen des Verdachts, am 11. März und 14. Juli 2005 an zwei Raubüberfällen gegen Tankstellen teilgenommen zu haben.
Am 25. September 2007 erhob die Staatsanwaltschaft Anklage wegen Vergewaltigung sowie wegen Teilnahme an einem Raubüberfall.
C. Am 19. Dezember 2007 wurde X._ wegen Drogenbesitzes verhaftet und am 21. Dezember 2007 aus der Haft entlassen. Am 21. Januar 2008 wurde er in diesem Zusammenhang erneut verhaftet. Mit Verfügung vom 4. Februar 2008 wurde ihm der vorzeitige Strafantritt bewilligt.
Am 14. März 2008 erhob die Staatsanwaltschaft Zusatzanklage betreffend des Vorwurfs der mehrfachen Widerhandlung gegen das Betäubungsmittelgesetz.
D. Am 1. April 2008 fand die Hauptverhandlung statt. Der Präsident der 2. Abteilung des Bezirksgerichts Zürich ordnete Beweisergänzungen an, u.a. die Einvernahme der Geschädigten A._ als Zeugin.
E. Am 2. April 2008 ersuchte X._ um Haftentlassung. Diesen Antrag wies der Haftrichter am Bezirksgericht Zürich mit Verfügung vom 10. April 2008 ab. Der Haftrichter bejahte Kollusionsgefahr aufgrund der vom Bezirksgericht angeordneten erneuten Befragung von A._.
Am 6. Mai 2008 hiess das Bundesgericht die dagegen gerichtete Beschwerde von X._ gut, weil keine Kollusionsgefahr vorliege. Es wies die Sache zur Prüfung allfälliger weiterer Haftgründe an den Haftrichter zurück.
F. Mit Verfügung vom 15. Mai 2008 wies der Haftrichter am Bezirksgericht Zürich das Gesuch um Entlassung aus dem vorzeitigen Strafantritt erneut ab, diesmal gestützt auf Fluchtgefahr und Wiederholungsgefahr.
Dagegen hat X._ am 20. Mai 2008 Beschwerde in Strafsachen an das Bundesgericht erhoben. Er beantragt, die angefochtene Verfügung sei aufzuheben und er sei aus der Haft zu entlassen. Eventualiter seien die zuständigen kantonalen Behörden anzuweisen, angemessene Auflagen oder Ersatzmassnahmen mit der Haftentlassung zu verbinden. Überdies beantragt er die Gewährung der unentgeltlichen Rechtspflege und Verbeiständung.
G. Die Staatsanwaltschaft und der Haftrichter haben auf eine Vernehmlassung verzichtet.

Erwägungen:
1. Da alle Sachurteilsvoraussetzungen vorliegen, ist auf die Beschwerde in Strafsachen grundsätzlich einzutreten. Nicht einzutreten ist allerdings auf die Rüge der Verletzung des Beschleunigungsgebots, die nicht (genügend) begründet wird (Art. 106 Abs. 2 i.V.m. <ref-law>).
2. Der Beschwerdeführer befindet sich im vorzeitigen Strafvollzug. Dieser setzt das Einverständnis der angeschuldigten Person sowie die Erwartung einer unbedingten Freiheitsstrafe oder einer stationären Massnahme voraus und darf den Zweck des Strafverfahrens nicht gefährden (<ref-law>/ZH). Für alle strafprozessualen Häftlinge, inklusive Gefangene im vorzeitigen Straf- und Massnahmenvollzug, gilt die Unschuldsvermutung (<ref-law>). Diese können sich auf die einschlägigen Verfahrensgarantien von <ref-law> berufen (<ref-ruling> E. 2 S. 275 mit Hinweisen) und jederzeit ein Begehren um Entlassung aus dem vorzeitigen Strafvollzug stellen (<ref-ruling> E. 1d S. 79 f.). Dieses Gesuch darf nur abgewiesen werden, wenn strafprozessuale Haftgründe fortdauern und die Dauer der Haft bzw. des Strafvollzugs nicht in die Nähe der konkret zu erwartenden Strafe gerückt ist (<ref-ruling> E. 3.4.2 S. 281 mit Hinweisen; <ref-ruling> E. 1d S. 80).
Voraussetzung für die Anordnung und Fortdauer von Untersuchungs- bzw. Sicherheitshaft (<ref-law>/ZH) ist nach zürcherischem Strafprozessrecht, dass gegen den Angeschuldigten der dringende Tatverdacht eines Vergehens oder Verbrechens besteht und zudem ein besonderer Haftgrund vorliegt, namentlich Kollusions-, Flucht-, oder Wiederholungsgefahr (<ref-law>/ZH). Die Untersuchungs- oder Sicherheitshaft ist durch mildere Massnahmen zu ersetzen, sofern sich der Haftzweck auch auf diese Weise erreichen lässt (§ 58 Abs. 4 i.V.m. § 72 f. StPO/ZH).
Im vorliegenden Fall ist ein dringender Tatverdacht im Hinblick auf die angeklagten Taten (Raub, Vergewaltigung, Vergehen gegen das Betäubungsmittelgesetz) zu bejahen. Streitig ist dagegen, ob der Haftgrund der Fluchtgefahr und/oder der Wiederholungsgefahr vorliegt und ob sich der Haftzweck durch mildere Massnahmen erreichen liesse.
3. Zu prüfen ist zunächst, ob Fluchtgefahr vorliegt.
3.1 Der Haftrichter bejahte Fluchtgefahr, weil der Beschwerdeführer im Falle der Verurteilung mit einer längeren Freiheitsstrafe und dem Widerruf seiner Niederlassungsbewilligung C durch das Migrationsamt rechnen müsse. Der Beschwerdeführer verfüge über keinen festen Wohnsitz in der Schweiz und es sei davon auszugehen, dass er familiäre und freundschaftliche Beziehungen in seinem Heimatland Tunesien unterhalte, in dem er im Sommer 2007 mehrere Monate verbracht habe. Der Haftrichter berücksichtigte, dass der Beschwerdeführer regelmässigen Kontakt zu Frau B._ in Schlieren unterhalte, die er zu heiraten gedenke, und bei der er im Falle der Haftentlassung unterkommen könne. Er räumte auch ein, dass sich der Beschwerdeführer seit Beginn der Untersuchung im Sommer 2005 nie der Strafverfolgung durch Flucht entzogen habe, sondern im Sommer 2007 sogar freiwillig in die Schweiz zurückgekehrt sei. Der Haftrichter ging aber davon aus, dass sich die Situation des Gesuchstellers seit der Durchführung der Hauptverhandlung in dem Sinne einschneidend verändert habe, als sich die Wahrscheinlichkeit einer Verurteilung zu einer unbedingten mehrjährigen Freiheitsstrafe mit anschliessendem Widerruf der Niederlassungsbewilligung konkretisiert habe und der Beschwerdeführer deshalb nun doch ernsthaft in Betracht ziehen könnte, sich ins Ausland abzusetzen. Demgegenüber gelte es sicherzustellen, dass die Untersuchung durch die noch notwendige Zeugenbefragung ergänzt und die Hauptverhandlung fortgesetzt werden könne.
3.2 Der Beschwerdeführer bestreitet, dass Fluchtgefahr bestehe. Er habe sich seit Einleitung des Strafverfahrens vor bald drei Jahren stets den Strafverfolgungsbehörden zur Verfügung gehalten und sich mehrfach freiwillig gemeldet, obschon er dabei insgesamt drei mal verhaftet worden sei. Auch von einem Aufenthalt in Tunesien sei er freiwillig in die Schweiz zurückgekehrt. Der Haftrichter habe deshalb auch in seiner Verfügung vom 23. Januar 2008 das Vorliegen von Fluchtgefahr verneint.
Es gebe keinen Grund, die Situation heute anders einzuschätzen, im Gegenteil: Die Beziehung des Beschwerdeführers zu Frau B._ und ihren zwei Kindern sei heute so intensiv geworden, dass sie die Gründung einer gemeinsamen Familie planten. Frau B._ könnte dem Beschwerdeführer nicht ins Ausland folgen, da ihr 10-jähriger Sohn ein Internat in Russikon besuche und ihr 16-jähriger Sohn in Frankreich studiere. Schon allein aus diesem Grund habe der Beschwerdeführer ein immenses Interesse, in der Schweiz zu bleiben.
Die Vorwürfe der Vergewaltigung und des Raubes und die deshalb drohende Verurteilung zu einer empfindlichen Strafe seien dem Beschwerdeführer schon lange bekannt, weshalb der Strafantrag der Staatsanwaltschaft von 5 1⁄2 Jahren Freiheitsstrafe in der Hauptverhandlung vom 1. April 2008 nicht überraschend gekommen sei. Ob die Niederlassungsbewilligung C infolge einer Verurteilung widerrufen werde, stehe noch nicht fest: Das Migrationsamt habe sich zu dieser Frage noch nicht geäussert und der Widerruf sei keineswegs zwingend.
3.3 Nach der Rechtsprechung des Bundesgerichts braucht es für die Annahme von Fluchtgefahr eine gewisse Wahrscheinlichkeit, dass sich der Angeschuldigte, wenn er in Freiheit wäre, der Strafverfolgung und dem Vollzug der Strafe durch Flucht entziehen würde. Die Schwere der drohenden Strafe darf als ein Indiz für Fluchtgefahr gewertet werden. Sie genügt jedoch für sich allein nicht, um den Haftgrund zu bejahen. Vielmehr müssen die konkreten Umstände des betreffenden Falles, insbesondere die gesamten Lebensverhältnisse des Angeschuldigten, in Betracht gezogen werden (<ref-ruling> E. 3a S. 62; <ref-ruling> E. 4a S. 70, je mit Hinweisen). Zu berücksichtigen sind beispielsweise die familiären und sozialen Bindungen des Häftlings, dessen berufliche Situation und Schulden sowie Kontakte ins Ausland und Ähnliches.
Im vorliegenden Fall hat der Haftrichter neben der drohenden Freiheitsstrafe die konkreten Umstände und die Lebensverhältnisse des Angeschuldigten gewürdigt, wobei er auch die gegen Fluchtgefahr sprechenden Indizien (bisheriges Verhalten im Strafverfahren; Beziehung zu Frau B._) berücksichtigt hat. In der Gesamtwürdigung bejahte er dennoch Fluchtgefahr, weil sich die Wahrscheinlichkeit einer Verurteilung zu einer unbedingten mehrjährigen Freiheitsstrafe mit anschliessendem Widerruf der Niederlassungsbewilligung inzwischen konkretisiert habe und der Beschwerdeführer deshalb ernsthaft in Betracht ziehen könnte, in sein Heimatland Tunesien zurückzukehren, in dem er noch über familiäre und freundschaftliche Kontakte verfüge.
Diese Würdigung ist nicht zu beanstanden: Zwar musste der Beschwerdeführer bereits aufgrund des Untersuchungsverfahrens und der Anklageerhebung mit der Möglichkeit einer empfindlichen Strafe rechnen. Dennoch stellt der Strafantrag der Staatsanwaltschaft an der Hauptverhandlung von 5 1⁄2 Jahren Freiheitsstrafe ein neues Element dar. Eine Verurteilung in dieser Grössenordnung würde mit grosser Wahrscheinlichkeit (wenn auch nicht zwingend) den Widerruf der Niederlassungsbewilligung nach sich ziehen (vgl. Art. 63 i.V.m. Art. 62 lit. b des Bundesgesetzes vom 16. Dezember 2005 über die Ausländerinnen und Ausländer [AuG; SR 142.20]). Muss der Beschwerdeführer deshalb ernsthaft damit rechnen, dass er die Schweiz ohnehin verlassen muss, erscheint die Versuchung, dies vor einer allfälligen Verurteilung und Vollstreckung einer langjährigen Freiheitsstrafe zu tun, gross. Hinzu kommt, dass der Beschwerdeführer, der seine Jugend in Tunesien verlebt hat, dort noch über familiäre und freundschaftliche Beziehungen verfügt. Es erscheint auch nicht von vornherein ausgeschlossen, dass ihm B._ und ihre Kinder nach Tunesien folgen könnten.
3.4 Nach dem Gesagten durfte der Haftrichter eine ernsthafte Gefahr bzw. Wahrscheinlichkeit der Flucht bejahen. Es kann daher offen bleiben, ob auch der Haftgrund der Wiederholungsgefahr besteht.
4. Bei der Prüfung der Verhältnismässigkeit der Haft ist zu berücksichtigen, dass das Strafverfahren weit fortgeschritten ist und in Kürze mit dem erstinstanzlichen Entscheid zu rechnen ist. Die Dauer der Untersuchungs- und Sicherheitshaft bzw. des vorzeitigen Strafvollzugs von insgesamt ca. 4 1⁄2 Monaten ist offensichtlich noch nicht in die Nähe der zu erwartenden Freiheitsstrafe gerückt.
Entgegen der Auffassung des Beschwerdeführers spielt es keine Rolle, dass die Untersuchungs- bzw. Sicherheitshaft im Januar 2008 wegen des Verdachts von Betäubungsmitteldelikten angeordnet wurde, die für das Strafmass nicht stark ins Gewicht fallen. Für die Verhältnismässigkeitsprüfung ist eine Gesamtbetrachtung unter Berücksichtigung aller Umstände des Falles erforderlich, wobei auf die heutige Situation des Beschwerdeführers, nach der Hauptverhandlung vom 1. April 2008, und die ihm insgesamt - für alle angeklagten Delikte zusammen - drohende Strafe abzustellen ist.
Die vom Beschwerdeführer vorgeschlagenen Ersatzmassnahmen (z.B. die Weisung, sich um Arbeit zu bemühen), sind in erster Linie auf den Haftgrund der Wiederholungsgefahr zugeschnitten. Die vom Beschwerdeführer vorgeschlagene Auflage, bei Frau B._ Wohnsitz zu nehmen und sich wöchentlich bei der Staatsanwaltschaft zu melden, erscheint nicht geeignet, den Beschwerdeführer an einem Untertauchen in oder einer Flucht aus der Schweiz zu hindern. Auch eine Pass- und Schriftensperre könnte das Risiko einer Flucht nicht genügend vermindern.
5. Der Beschwerdeführer rügt in diesem Zusammenhang eine Verletzung des rechtlichen Gehörs, weil sich der Haftrichter mit dem Eventualantrag auf Haftentlassung unter angemessenen Auflagen resp. Ersatzmassnahmen und den vom Beschwerdeführer konkret vorgeschlagenen Massnahmen nicht auseinandergesetzt habe.
Der Haftrichter muss stets prüfen, ob die Untersuchungs- oder Sicherheitshaft durch mildere Massnahmen ersetzt werden kann (§ 58 Abs. 4 i.V.m. § 72 f. StPO/ZH; dies ergibt sich auch aus dem Verhältnismässigkeitsprinzip; vgl. dazu <ref-ruling> E. 3.3 S. 279). Dies bedeutet aber nicht, dass der Haftentscheid immer Erwägungen zu dieser Frage enthalten muss, darf sich doch die Begründung auf die für den Entscheid wesentlichen Punkte beschränken (<ref-ruling> E. 4.1 mit Hinweisen). Eine Verletzung der Begründungspflicht und des rechtlichen Gehörs liegt deshalb nur vor, wenn sich Ausführungen zu Ersatzmassnahmen aufgrund der konkreten Umstände des Falles aufdrängen, insbesondere wenn konkrete geeignete Massnahmen vom Beschwerdeführer vorgeschlagen worden sind (vgl. <ref-ruling> E. 3.3.3 S. 280).
Im vorliegenden Fall hatte die Rechtsvertreterin des Beschwerdeführers in ihrer Eingabe vom 15. Mai 2008 als mildere Massnahme im Zusammenhang mit der Fluchtgefahr lediglich eine Meldepflicht bei der Staatsanwaltschaft vorgeschlagen. Da diese Massnahme nach dem oben Gesagten nicht geeignet war, den Haftzweck zu ersetzen, war der Haftrichter nicht gehalten, sich speziell dazu zu äussern.
6. Nach dem Gesagten erweist sich die Beschwerde als unbegründet und ist abzuweisen, soweit darauf einzutreten ist.
Da die Voraussetzungen für die Gewährung der unentgeltlichen Rechtspflege und Verbeiständung vorliegen (Art. 64 Abs. 1 und 2 BGG), sind keine Kosten zu erheben und dem Beschwerdeführer ist Rechtsanwältin Antigone Schobinger als unentgeltliche Rechtsvertreterin zu bestellen.

Demnach erkennt das Bundesgericht:
1. Die Beschwerde wird abgewiesen, soweit darauf einzutreten ist.
2. Dem Beschwerdeführer wird die unentgeltliche Rechtspflege gewährt.
2.1 Es werden keine Kosten erhoben.
2.2 Rechtsanwältin Antigone Schobinger, wird als amtliche Vertreterin des Beschwerdeführers bestellt, und es wird ihr für das bundesgerichtliche Verfahren aus der Bundesgerichtskasse ein Honorar von Fr. 1'500.-- ausgerichtet.
3. Dieses Urteil wird dem Beschwerdeführer, der Staatsanwaltschaft Limmattal/Albis, Zweigstelle Dietikon, und dem Bezirksgericht Zürich, Haftrichter, schriftlich mitgeteilt.
Lausanne, 30. Mai 2008
Im Namen der I. öffentlich-rechtlichen Abteilung
des Schweizerischen Bundesgerichts
Der Präsident: Die Gerichtsschreiberin:
Féraud Gerber